1use bytes::Bytes;
21use half::f16;
22use std::cmp::Ordering;
23use std::fmt;
24use std::mem;
25use std::ops::{Deref, DerefMut};
26use std::str::from_utf8;
27
28use crate::basic::Type;
29use crate::column::reader::{ColumnReader, ColumnReaderImpl};
30use crate::column::writer::{ColumnWriter, ColumnWriterImpl};
31use crate::errors::{ParquetError, Result};
32use crate::util::bit_util::FromBytes;
33
34#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
37pub struct Int96 {
38 value: [u32; 3],
39}
40
41const JULIAN_DAY_OF_EPOCH: i64 = 2_440_588;
42
43const SECONDS_IN_DAY: i64 = 86_400;
45const MILLISECONDS: i64 = 1_000;
47const MICROSECONDS: i64 = 1_000_000;
49const NANOSECONDS: i64 = 1_000_000_000;
51
52const MILLISECONDS_IN_DAY: i64 = SECONDS_IN_DAY * MILLISECONDS;
54const MICROSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * MICROSECONDS;
56const NANOSECONDS_IN_DAY: i64 = SECONDS_IN_DAY * NANOSECONDS;
58
59impl Int96 {
60 pub fn new() -> Self {
62 Self { value: [0; 3] }
63 }
64
65 #[inline]
67 pub fn data(&self) -> &[u32] {
68 &self.value
69 }
70
71 #[inline]
73 pub fn set_data(&mut self, elem0: u32, elem1: u32, elem2: u32) {
74 self.value = [elem0, elem1, elem2];
75 }
76
77 #[inline]
81 pub fn to_seconds(&self) -> i64 {
82 let (day, nanos) = self.data_as_days_and_nanos();
83 (day as i64 - JULIAN_DAY_OF_EPOCH)
84 .wrapping_mul(SECONDS_IN_DAY)
85 .wrapping_add(nanos / 1_000_000_000)
86 }
87
88 #[inline]
92 pub fn to_millis(&self) -> i64 {
93 let (day, nanos) = self.data_as_days_and_nanos();
94 (day as i64 - JULIAN_DAY_OF_EPOCH)
95 .wrapping_mul(MILLISECONDS_IN_DAY)
96 .wrapping_add(nanos / 1_000_000)
97 }
98
99 #[inline]
103 pub fn to_micros(&self) -> i64 {
104 let (day, nanos) = self.data_as_days_and_nanos();
105 (day as i64 - JULIAN_DAY_OF_EPOCH)
106 .wrapping_mul(MICROSECONDS_IN_DAY)
107 .wrapping_add(nanos / 1_000)
108 }
109
110 #[inline]
114 pub fn to_nanos(&self) -> i64 {
115 let (day, nanos) = self.data_as_days_and_nanos();
116 (day as i64 - JULIAN_DAY_OF_EPOCH)
117 .wrapping_mul(NANOSECONDS_IN_DAY)
118 .wrapping_add(nanos)
119 }
120
121 #[inline]
122 fn get_days(&self) -> i32 {
123 self.data()[2] as i32
124 }
125
126 #[inline]
127 fn get_nanos(&self) -> i64 {
128 ((self.data()[1] as i64) << 32) + self.data()[0] as i64
129 }
130
131 #[inline]
132 fn data_as_days_and_nanos(&self) -> (i32, i64) {
133 (self.get_days(), self.get_nanos())
134 }
135}
136
137impl PartialOrd for Int96 {
138 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
139 Some(self.cmp(other))
140 }
141}
142
143impl Ord for Int96 {
144 fn cmp(&self, other: &Self) -> Ordering {
153 match self.get_days().cmp(&other.get_days()) {
154 Ordering::Equal => self.get_nanos().cmp(&other.get_nanos()),
155 ord => ord,
156 }
157 }
158}
159impl From<Vec<u32>> for Int96 {
160 fn from(buf: Vec<u32>) -> Self {
161 assert_eq!(buf.len(), 3);
162 let mut result = Self::new();
163 result.set_data(buf[0], buf[1], buf[2]);
164 result
165 }
166}
167
168impl fmt::Display for Int96 {
169 #[cold]
170 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
171 write!(f, "{:?}", self.data())
172 }
173}
174
175#[derive(Clone, Default)]
178pub struct ByteArray {
179 data: Option<Bytes>,
180}
181
182impl std::fmt::Debug for ByteArray {
184 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
185 let mut debug_struct = f.debug_struct("ByteArray");
186 match self.as_utf8() {
187 Ok(s) => debug_struct.field("data", &s),
188 Err(_) => debug_struct.field("data", &self.data),
189 };
190 debug_struct.finish()
191 }
192}
193
194impl PartialOrd for ByteArray {
195 fn partial_cmp(&self, other: &ByteArray) -> Option<Ordering> {
196 match (&self.data, &other.data) {
201 (None, None) => Some(Ordering::Equal),
202 (None, Some(_)) => Some(Ordering::Less),
203 (Some(_), None) => Some(Ordering::Greater),
204 (Some(self_data), Some(other_data)) => {
205 self_data.partial_cmp(&other_data)
207 }
208 }
209 }
210}
211
212impl ByteArray {
213 #[inline]
215 pub fn new() -> Self {
216 ByteArray { data: None }
217 }
218
219 #[inline]
225 pub fn len(&self) -> usize {
226 assert!(self.data.is_some());
227 self.data.as_ref().unwrap().len()
228 }
229
230 #[inline]
236 pub fn is_empty(&self) -> bool {
237 self.len() == 0
238 }
239
240 #[inline]
246 pub fn data(&self) -> &[u8] {
247 self.data
248 .as_ref()
249 .expect("set_data should have been called")
250 .as_ref()
251 }
252
253 #[inline]
255 pub fn set_data(&mut self, data: Bytes) {
256 self.data = Some(data);
257 }
258
259 #[inline]
265 pub fn slice(&self, start: usize, len: usize) -> Self {
266 Self::from(
267 self.data
268 .as_ref()
269 .expect("set_data should have been called")
270 .slice(start..start + len),
271 )
272 }
273
274 pub fn as_utf8(&self) -> Result<&str> {
276 let bytes = self
277 .data
278 .as_ref()
279 .map(|ptr| ptr.as_ref())
280 .ok_or_else(|| general_err!("Can't convert empty byte array to utf8"))?;
281 from_utf8(bytes).map_err(|e| e.into())
282 }
283}
284
285impl From<Vec<u8>> for ByteArray {
286 fn from(buf: Vec<u8>) -> ByteArray {
287 Self {
288 data: Some(buf.into()),
289 }
290 }
291}
292
293impl<'a> From<&'a [u8]> for ByteArray {
294 fn from(b: &'a [u8]) -> ByteArray {
295 let mut v = Vec::new();
296 v.extend_from_slice(b);
297 Self {
298 data: Some(v.into()),
299 }
300 }
301}
302
303impl<'a> From<&'a str> for ByteArray {
304 fn from(s: &'a str) -> ByteArray {
305 let mut v = Vec::new();
306 v.extend_from_slice(s.as_bytes());
307 Self {
308 data: Some(v.into()),
309 }
310 }
311}
312
313impl From<Bytes> for ByteArray {
314 fn from(value: Bytes) -> Self {
315 Self { data: Some(value) }
316 }
317}
318
319impl From<f16> for ByteArray {
320 fn from(value: f16) -> Self {
321 Self::from(value.to_le_bytes().as_slice())
322 }
323}
324
325impl PartialEq for ByteArray {
326 fn eq(&self, other: &ByteArray) -> bool {
327 match (&self.data, &other.data) {
328 (Some(d1), Some(d2)) => d1.as_ref() == d2.as_ref(),
329 (None, None) => true,
330 _ => false,
331 }
332 }
333}
334
335impl fmt::Display for ByteArray {
336 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
337 write!(f, "{:?}", self.data())
338 }
339}
340
341#[repr(transparent)]
356#[derive(Clone, Debug, Default)]
357pub struct FixedLenByteArray(ByteArray);
358
359impl PartialEq for FixedLenByteArray {
360 fn eq(&self, other: &FixedLenByteArray) -> bool {
361 self.0.eq(&other.0)
362 }
363}
364
365impl PartialEq<ByteArray> for FixedLenByteArray {
366 fn eq(&self, other: &ByteArray) -> bool {
367 self.0.eq(other)
368 }
369}
370
371impl PartialEq<FixedLenByteArray> for ByteArray {
372 fn eq(&self, other: &FixedLenByteArray) -> bool {
373 self.eq(&other.0)
374 }
375}
376
377impl fmt::Display for FixedLenByteArray {
378 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
379 self.0.fmt(f)
380 }
381}
382
383impl PartialOrd for FixedLenByteArray {
384 fn partial_cmp(&self, other: &FixedLenByteArray) -> Option<Ordering> {
385 self.0.partial_cmp(&other.0)
386 }
387}
388
389impl PartialOrd<FixedLenByteArray> for ByteArray {
390 fn partial_cmp(&self, other: &FixedLenByteArray) -> Option<Ordering> {
391 self.partial_cmp(&other.0)
392 }
393}
394
395impl PartialOrd<ByteArray> for FixedLenByteArray {
396 fn partial_cmp(&self, other: &ByteArray) -> Option<Ordering> {
397 self.0.partial_cmp(other)
398 }
399}
400
401impl Deref for FixedLenByteArray {
402 type Target = ByteArray;
403
404 fn deref(&self) -> &Self::Target {
405 &self.0
406 }
407}
408
409impl DerefMut for FixedLenByteArray {
410 fn deref_mut(&mut self) -> &mut Self::Target {
411 &mut self.0
412 }
413}
414
415impl From<ByteArray> for FixedLenByteArray {
416 fn from(other: ByteArray) -> Self {
417 Self(other)
418 }
419}
420
421impl From<Vec<u8>> for FixedLenByteArray {
422 fn from(buf: Vec<u8>) -> FixedLenByteArray {
423 FixedLenByteArray(ByteArray::from(buf))
424 }
425}
426
427impl From<FixedLenByteArray> for ByteArray {
428 fn from(other: FixedLenByteArray) -> Self {
429 other.0
430 }
431}
432
433#[derive(Clone, Debug)]
439pub enum Decimal {
440 Int32 {
442 value: [u8; 4],
444 precision: i32,
446 scale: i32,
448 },
449 Int64 {
451 value: [u8; 8],
453 precision: i32,
455 scale: i32,
457 },
458 Bytes {
460 value: ByteArray,
462 precision: i32,
464 scale: i32,
466 },
467}
468
469impl Decimal {
470 pub fn from_i32(value: i32, precision: i32, scale: i32) -> Self {
472 let bytes = value.to_be_bytes();
473 Decimal::Int32 {
474 value: bytes,
475 precision,
476 scale,
477 }
478 }
479
480 pub fn from_i64(value: i64, precision: i32, scale: i32) -> Self {
482 let bytes = value.to_be_bytes();
483 Decimal::Int64 {
484 value: bytes,
485 precision,
486 scale,
487 }
488 }
489
490 pub fn from_bytes(value: ByteArray, precision: i32, scale: i32) -> Self {
492 Decimal::Bytes {
493 value,
494 precision,
495 scale,
496 }
497 }
498
499 pub fn data(&self) -> &[u8] {
501 match *self {
502 Decimal::Int32 { ref value, .. } => value,
503 Decimal::Int64 { ref value, .. } => value,
504 Decimal::Bytes { ref value, .. } => value.data(),
505 }
506 }
507
508 pub fn precision(&self) -> i32 {
510 match *self {
511 Decimal::Int32 { precision, .. } => precision,
512 Decimal::Int64 { precision, .. } => precision,
513 Decimal::Bytes { precision, .. } => precision,
514 }
515 }
516
517 pub fn scale(&self) -> i32 {
519 match *self {
520 Decimal::Int32 { scale, .. } => scale,
521 Decimal::Int64 { scale, .. } => scale,
522 Decimal::Bytes { scale, .. } => scale,
523 }
524 }
525}
526
527impl Default for Decimal {
528 fn default() -> Self {
529 Self::from_i32(0, 0, 0)
530 }
531}
532
533impl PartialEq for Decimal {
534 fn eq(&self, other: &Decimal) -> bool {
535 self.precision() == other.precision()
536 && self.scale() == other.scale()
537 && self.data() == other.data()
538 }
539}
540
541pub trait AsBytes {
543 fn as_bytes(&self) -> &[u8];
545}
546
547pub trait SliceAsBytes: Sized {
549 fn slice_as_bytes(self_: &[Self]) -> &[u8];
551 unsafe fn slice_as_bytes_mut(self_: &mut [Self]) -> &mut [u8];
557}
558
559impl AsBytes for [u8] {
560 fn as_bytes(&self) -> &[u8] {
561 self
562 }
563}
564
565macro_rules! gen_as_bytes {
566 ($source_ty:ident) => {
567 impl AsBytes for $source_ty {
568 fn as_bytes(&self) -> &[u8] {
569 unsafe {
572 std::slice::from_raw_parts(
573 std::ptr::from_ref::<$source_ty>(self).cast::<u8>(),
574 std::mem::size_of::<$source_ty>(),
575 )
576 }
577 }
578 }
579
580 impl SliceAsBytes for $source_ty {
581 #[inline]
582 fn slice_as_bytes(self_: &[Self]) -> &[u8] {
583 unsafe {
586 std::slice::from_raw_parts(
587 self_.as_ptr().cast::<u8>(),
588 std::mem::size_of_val(self_),
589 )
590 }
591 }
592
593 #[inline]
594 unsafe fn slice_as_bytes_mut(self_: &mut [Self]) -> &mut [u8] {
595 unsafe {
599 std::slice::from_raw_parts_mut(
600 self_.as_mut_ptr().cast::<u8>(),
601 std::mem::size_of_val(self_),
602 )
603 }
604 }
605 }
606 };
607}
608
609gen_as_bytes!(i8);
610gen_as_bytes!(i16);
611gen_as_bytes!(i32);
612gen_as_bytes!(i64);
613gen_as_bytes!(u8);
614gen_as_bytes!(u16);
615gen_as_bytes!(u32);
616gen_as_bytes!(u64);
617gen_as_bytes!(f32);
618gen_as_bytes!(f64);
619
620macro_rules! unimplemented_slice_as_bytes {
621 ($ty: ty) => {
622 impl SliceAsBytes for $ty {
623 fn slice_as_bytes(_self: &[Self]) -> &[u8] {
624 unimplemented!()
625 }
626
627 unsafe fn slice_as_bytes_mut(_self: &mut [Self]) -> &mut [u8] {
628 unimplemented!()
629 }
630 }
631 };
632}
633
634unimplemented_slice_as_bytes!(Int96);
636unimplemented_slice_as_bytes!(bool);
637unimplemented_slice_as_bytes!(ByteArray);
638unimplemented_slice_as_bytes!(FixedLenByteArray);
639
640impl AsBytes for bool {
641 fn as_bytes(&self) -> &[u8] {
642 unsafe { std::slice::from_raw_parts(std::ptr::from_ref::<bool>(self).cast::<u8>(), 1) }
645 }
646}
647
648impl AsBytes for Int96 {
649 fn as_bytes(&self) -> &[u8] {
650 unsafe {
652 std::slice::from_raw_parts(std::ptr::from_ref::<[u32]>(self.data()).cast::<u8>(), 12)
653 }
654 }
655}
656
657impl AsBytes for ByteArray {
658 fn as_bytes(&self) -> &[u8] {
659 self.data()
660 }
661}
662
663impl AsBytes for FixedLenByteArray {
664 fn as_bytes(&self) -> &[u8] {
665 self.data()
666 }
667}
668
669impl AsBytes for Decimal {
670 fn as_bytes(&self) -> &[u8] {
671 self.data()
672 }
673}
674
675impl AsBytes for Vec<u8> {
676 fn as_bytes(&self) -> &[u8] {
677 self.as_slice()
678 }
679}
680
681impl AsBytes for &str {
682 fn as_bytes(&self) -> &[u8] {
683 (self as &str).as_bytes()
684 }
685}
686
687impl AsBytes for str {
688 fn as_bytes(&self) -> &[u8] {
689 (self as &str).as_bytes()
690 }
691}
692
693pub(crate) mod private {
694 use bytes::Bytes;
695
696 use crate::encodings::decoding::PlainDecoderDetails;
697 use crate::util::bit_util::{BitReader, BitWriter, read_num_bytes};
698
699 use super::{ParquetError, Result, SliceAsBytes};
700 use crate::basic::Type;
701 use crate::file::metadata::HeapSize;
702
703 pub trait ParquetValueType:
709 PartialEq
710 + std::fmt::Debug
711 + std::fmt::Display
712 + Default
713 + Clone
714 + super::AsBytes
715 + super::FromBytes
716 + SliceAsBytes
717 + PartialOrd
718 + Send
719 + HeapSize
720 + crate::encodings::decoding::private::GetDecoder
721 + crate::encodings::encoding::private::GetEncoder
722 + crate::file::statistics::private::MakeStatistics
723 {
724 const PHYSICAL_TYPE: Type;
725
726 fn encode<W: std::io::Write>(
728 values: &[Self],
729 writer: &mut W,
730 bit_writer: &mut BitWriter,
731 ) -> Result<()>;
732
733 fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize);
735
736 fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result<usize>;
738
739 fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result<usize>;
740
741 fn dict_encoding_size(&self) -> (usize, usize) {
743 (std::mem::size_of::<Self>(), 1)
744 }
745
746 fn variable_length_bytes(_: &[Self]) -> Option<i64> {
750 None
751 }
752
753 fn as_i64(&self) -> Result<i64> {
758 Err(general_err!("Type cannot be converted to i64"))
759 }
760
761 fn as_u64(&self) -> Result<u64> {
766 self.as_i64()
767 .map_err(|_| general_err!("Type cannot be converted to u64"))
768 .map(|x| x as u64)
769 }
770
771 fn as_any(&self) -> &dyn std::any::Any;
773
774 fn as_mut_any(&mut self) -> &mut dyn std::any::Any;
776
777 fn set_from_bytes(&mut self, _data: Bytes) {
781 unimplemented!();
782 }
783 }
784
785 impl ParquetValueType for bool {
786 const PHYSICAL_TYPE: Type = Type::BOOLEAN;
787
788 #[inline]
789 fn encode<W: std::io::Write>(
790 values: &[Self],
791 _: &mut W,
792 bit_writer: &mut BitWriter,
793 ) -> Result<()> {
794 bit_writer.put_batch(values, 1);
795 Ok(())
796 }
797
798 #[inline]
799 fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) {
800 decoder.bit_reader.replace(BitReader::new(data));
801 decoder.num_values = num_values;
802 }
803
804 #[inline]
805 fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result<usize> {
806 let bit_reader = decoder.bit_reader.as_mut().unwrap();
807 let num_values = std::cmp::min(buffer.len(), decoder.num_values);
808 let values_read = bit_reader.get_batch(&mut buffer[..num_values], 1);
809 decoder.num_values -= values_read;
810 Ok(values_read)
811 }
812
813 fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result<usize> {
814 let bit_reader = decoder.bit_reader.as_mut().unwrap();
815 let num_values = std::cmp::min(num_values, decoder.num_values);
816 let values_read = bit_reader.skip(num_values, 1);
817 decoder.num_values -= values_read;
818 Ok(values_read)
819 }
820
821 #[inline]
822 fn as_i64(&self) -> Result<i64> {
823 Ok(*self as i64)
824 }
825
826 #[inline]
827 fn as_any(&self) -> &dyn std::any::Any {
828 self
829 }
830
831 #[inline]
832 fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
833 self
834 }
835 }
836
837 macro_rules! impl_from_raw {
838 ($ty: ty, $physical_ty: expr, $self: ident => $as_i64: block) => {
839 impl ParquetValueType for $ty {
840 const PHYSICAL_TYPE: Type = $physical_ty;
841
842 #[inline]
843 fn encode<W: std::io::Write>(values: &[Self], writer: &mut W, _: &mut BitWriter) -> Result<()> {
844 let raw = unsafe {
846 std::slice::from_raw_parts(
847 values.as_ptr().cast::<u8>(),
848 std::mem::size_of_val(values),
849 )
850 };
851 writer.write_all(raw)?;
852
853 Ok(())
854 }
855
856 #[inline]
857 fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) {
858 decoder.data.replace(data);
859 decoder.start = 0;
860 decoder.num_values = num_values;
861 }
862
863 #[inline]
864 fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result<usize> {
865 let data = decoder.data.as_ref().expect("set_data should have been called");
866 let num_values = std::cmp::min(buffer.len(), decoder.num_values);
867 let bytes_left = data.len() - decoder.start;
868 let bytes_to_decode = std::mem::size_of::<Self>() * num_values;
869
870 if bytes_left < bytes_to_decode {
871 return Err(eof_err!("Not enough bytes to decode"));
872 }
873
874 {
875 let raw_buffer = &mut unsafe { Self::slice_as_bytes_mut(buffer) }[..bytes_to_decode];
878 raw_buffer.copy_from_slice(data.slice(
879 decoder.start..decoder.start + bytes_to_decode
880 ).as_ref());
881 };
882 decoder.start += bytes_to_decode;
883 decoder.num_values -= num_values;
884
885 Ok(num_values)
886 }
887
888 #[inline]
889 fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result<usize> {
890 let data = decoder.data.as_ref().expect("set_data should have been called");
891 let num_values = num_values.min(decoder.num_values);
892 let bytes_left = data.len() - decoder.start;
893 let bytes_to_skip = std::mem::size_of::<Self>() * num_values;
894
895 if bytes_left < bytes_to_skip {
896 return Err(eof_err!("Not enough bytes to skip"));
897 }
898
899 decoder.start += bytes_to_skip;
900 decoder.num_values -= num_values;
901
902 Ok(num_values)
903 }
904
905 #[inline]
906 fn as_i64(&$self) -> Result<i64> {
907 $as_i64
908 }
909
910 #[inline]
911 fn as_any(&self) -> &dyn std::any::Any {
912 self
913 }
914
915 #[inline]
916 fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
917 self
918 }
919 }
920 }
921 }
922
923 impl_from_raw!(i32, Type::INT32, self => { Ok(*self as i64) });
924 impl_from_raw!(i64, Type::INT64, self => { Ok(*self) });
925 impl_from_raw!(f32, Type::FLOAT, self => { Err(general_err!("Type cannot be converted to i64")) });
926 impl_from_raw!(f64, Type::DOUBLE, self => { Err(general_err!("Type cannot be converted to i64")) });
927
928 impl ParquetValueType for super::Int96 {
929 const PHYSICAL_TYPE: Type = Type::INT96;
930
931 #[inline]
932 fn encode<W: std::io::Write>(
933 values: &[Self],
934 writer: &mut W,
935 _: &mut BitWriter,
936 ) -> Result<()> {
937 for value in values {
938 let raw = SliceAsBytes::slice_as_bytes(value.data());
939 writer.write_all(raw)?;
940 }
941 Ok(())
942 }
943
944 #[inline]
945 fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) {
946 decoder.data.replace(data);
947 decoder.start = 0;
948 decoder.num_values = num_values;
949 }
950
951 #[inline]
952 fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result<usize> {
953 let data = decoder
955 .data
956 .as_ref()
957 .expect("set_data should have been called");
958 let num_values = std::cmp::min(buffer.len(), decoder.num_values);
959 let bytes_left = data.len() - decoder.start;
960 let bytes_to_decode = 12 * num_values;
961
962 if bytes_left < bytes_to_decode {
963 return Err(eof_err!("Not enough bytes to decode"));
964 }
965
966 let data_range = data.slice(decoder.start..decoder.start + bytes_to_decode);
967 let bytes: &[u8] = &data_range;
968 decoder.start += bytes_to_decode;
969
970 let mut pos = 0; for item in buffer.iter_mut().take(num_values) {
972 let elem0 = u32::from_le_bytes(bytes[pos..pos + 4].try_into().unwrap());
973 let elem1 = u32::from_le_bytes(bytes[pos + 4..pos + 8].try_into().unwrap());
974 let elem2 = u32::from_le_bytes(bytes[pos + 8..pos + 12].try_into().unwrap());
975
976 item.set_data(elem0, elem1, elem2);
977 pos += 12;
978 }
979 decoder.num_values -= num_values;
980
981 Ok(num_values)
982 }
983
984 fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result<usize> {
985 let data = decoder
986 .data
987 .as_ref()
988 .expect("set_data should have been called");
989 let num_values = std::cmp::min(num_values, decoder.num_values);
990 let bytes_left = data.len() - decoder.start;
991 let bytes_to_skip = 12 * num_values;
992
993 if bytes_left < bytes_to_skip {
994 return Err(eof_err!("Not enough bytes to skip"));
995 }
996 decoder.start += bytes_to_skip;
997 decoder.num_values -= num_values;
998
999 Ok(num_values)
1000 }
1001
1002 #[inline]
1003 fn as_any(&self) -> &dyn std::any::Any {
1004 self
1005 }
1006
1007 #[inline]
1008 fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
1009 self
1010 }
1011 }
1012
1013 impl HeapSize for super::Int96 {
1014 fn heap_size(&self) -> usize {
1015 0 }
1017 }
1018
1019 impl ParquetValueType for super::ByteArray {
1020 const PHYSICAL_TYPE: Type = Type::BYTE_ARRAY;
1021
1022 #[inline]
1023 fn encode<W: std::io::Write>(
1024 values: &[Self],
1025 writer: &mut W,
1026 _: &mut BitWriter,
1027 ) -> Result<()> {
1028 for value in values {
1029 let len: u32 = value.len().try_into().unwrap();
1030 writer.write_all(&len.to_ne_bytes())?;
1031 let raw = value.data();
1032 writer.write_all(raw)?;
1033 }
1034 Ok(())
1035 }
1036
1037 #[inline]
1038 fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) {
1039 decoder.data.replace(data);
1040 decoder.start = 0;
1041 decoder.num_values = num_values;
1042 }
1043
1044 #[inline]
1045 fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result<usize> {
1046 let data = decoder
1047 .data
1048 .as_mut()
1049 .expect("set_data should have been called");
1050 let num_values = std::cmp::min(buffer.len(), decoder.num_values);
1051 for val_array in buffer.iter_mut().take(num_values) {
1052 let len: usize =
1053 read_num_bytes::<u32>(4, data.slice(decoder.start..).as_ref()) as usize;
1054 decoder.start += std::mem::size_of::<u32>();
1055
1056 if data.len() < decoder.start + len {
1057 return Err(eof_err!("Not enough bytes to decode"));
1058 }
1059
1060 val_array.set_data(data.slice(decoder.start..decoder.start + len));
1061 decoder.start += len;
1062 }
1063 decoder.num_values -= num_values;
1064
1065 Ok(num_values)
1066 }
1067
1068 fn variable_length_bytes(values: &[Self]) -> Option<i64> {
1069 Some(values.iter().map(|x| x.len() as i64).sum())
1070 }
1071
1072 fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result<usize> {
1073 let data = decoder
1074 .data
1075 .as_mut()
1076 .expect("set_data should have been called");
1077 let num_values = num_values.min(decoder.num_values);
1078
1079 for _ in 0..num_values {
1080 let len: usize =
1081 read_num_bytes::<u32>(4, data.slice(decoder.start..).as_ref()) as usize;
1082 decoder.start += std::mem::size_of::<u32>() + len;
1083 }
1084 decoder.num_values -= num_values;
1085
1086 Ok(num_values)
1087 }
1088
1089 #[inline]
1090 fn dict_encoding_size(&self) -> (usize, usize) {
1091 (std::mem::size_of::<u32>(), self.len())
1092 }
1093
1094 #[inline]
1095 fn as_any(&self) -> &dyn std::any::Any {
1096 self
1097 }
1098
1099 #[inline]
1100 fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
1101 self
1102 }
1103
1104 #[inline]
1105 fn set_from_bytes(&mut self, data: Bytes) {
1106 self.set_data(data);
1107 }
1108 }
1109
1110 impl HeapSize for super::ByteArray {
1111 fn heap_size(&self) -> usize {
1112 self.data.as_ref().map(|data| data.len()).unwrap_or(0)
1116 }
1117 }
1118
1119 impl ParquetValueType for super::FixedLenByteArray {
1120 const PHYSICAL_TYPE: Type = Type::FIXED_LEN_BYTE_ARRAY;
1121
1122 #[inline]
1123 fn encode<W: std::io::Write>(
1124 values: &[Self],
1125 writer: &mut W,
1126 _: &mut BitWriter,
1127 ) -> Result<()> {
1128 for value in values {
1129 let raw = value.data();
1130 writer.write_all(raw)?;
1131 }
1132 Ok(())
1133 }
1134
1135 #[inline]
1136 fn set_data(decoder: &mut PlainDecoderDetails, data: Bytes, num_values: usize) {
1137 decoder.data.replace(data);
1138 decoder.start = 0;
1139 decoder.num_values = num_values;
1140 }
1141
1142 #[inline]
1143 fn decode(buffer: &mut [Self], decoder: &mut PlainDecoderDetails) -> Result<usize> {
1144 assert!(decoder.type_length > 0);
1145
1146 let data = decoder
1147 .data
1148 .as_mut()
1149 .expect("set_data should have been called");
1150 let num_values = std::cmp::min(buffer.len(), decoder.num_values);
1151
1152 for item in buffer.iter_mut().take(num_values) {
1153 let len = decoder.type_length as usize;
1154
1155 if data.len() < decoder.start + len {
1156 return Err(eof_err!("Not enough bytes to decode"));
1157 }
1158
1159 item.set_data(data.slice(decoder.start..decoder.start + len));
1160 decoder.start += len;
1161 }
1162 decoder.num_values -= num_values;
1163
1164 Ok(num_values)
1165 }
1166
1167 fn skip(decoder: &mut PlainDecoderDetails, num_values: usize) -> Result<usize> {
1168 assert!(decoder.type_length > 0);
1169
1170 let data = decoder
1171 .data
1172 .as_mut()
1173 .expect("set_data should have been called");
1174 let num_values = std::cmp::min(num_values, decoder.num_values);
1175 for _ in 0..num_values {
1176 let len = decoder.type_length as usize;
1177
1178 if data.len() < decoder.start + len {
1179 return Err(eof_err!("Not enough bytes to skip"));
1180 }
1181
1182 decoder.start += len;
1183 }
1184 decoder.num_values -= num_values;
1185
1186 Ok(num_values)
1187 }
1188
1189 #[inline]
1190 fn dict_encoding_size(&self) -> (usize, usize) {
1191 (std::mem::size_of::<u32>(), self.len())
1192 }
1193
1194 #[inline]
1195 fn as_any(&self) -> &dyn std::any::Any {
1196 self
1197 }
1198
1199 #[inline]
1200 fn as_mut_any(&mut self) -> &mut dyn std::any::Any {
1201 self
1202 }
1203
1204 #[inline]
1205 fn set_from_bytes(&mut self, data: Bytes) {
1206 self.set_data(data);
1207 }
1208 }
1209
1210 impl HeapSize for super::FixedLenByteArray {
1211 fn heap_size(&self) -> usize {
1212 self.0.heap_size()
1213 }
1214 }
1215}
1216
1217pub trait DataType: 'static + Send {
1220 type T: private::ParquetValueType;
1222
1223 fn get_physical_type() -> Type {
1225 <Self::T as private::ParquetValueType>::PHYSICAL_TYPE
1226 }
1227
1228 fn get_type_size() -> usize;
1230
1231 fn get_column_reader(column_writer: ColumnReader) -> Option<ColumnReaderImpl<Self>>
1233 where
1234 Self: Sized;
1235
1236 fn get_column_writer(column_writer: ColumnWriter<'_>) -> Option<ColumnWriterImpl<'_, Self>>
1238 where
1239 Self: Sized;
1240
1241 fn get_column_writer_ref<'a, 'b: 'a>(
1243 column_writer: &'b ColumnWriter<'a>,
1244 ) -> Option<&'b ColumnWriterImpl<'a, Self>>
1245 where
1246 Self: Sized;
1247
1248 fn get_column_writer_mut<'a, 'b: 'a>(
1250 column_writer: &'a mut ColumnWriter<'b>,
1251 ) -> Option<&'a mut ColumnWriterImpl<'b, Self>>
1252 where
1253 Self: Sized;
1254}
1255
1256macro_rules! make_type {
1257 ($name:ident, $reader_ident: ident, $writer_ident: ident, $native_ty:ty, $size:expr) => {
1258 #[doc = concat!("Parquet physical type: ", stringify!($name))]
1259 #[derive(Clone)]
1260 pub struct $name {}
1261
1262 impl DataType for $name {
1263 type T = $native_ty;
1264
1265 fn get_type_size() -> usize {
1266 $size
1267 }
1268
1269 fn get_column_reader(column_reader: ColumnReader) -> Option<ColumnReaderImpl<Self>> {
1270 match column_reader {
1271 ColumnReader::$reader_ident(w) => Some(w),
1272 _ => None,
1273 }
1274 }
1275
1276 fn get_column_writer(
1277 column_writer: ColumnWriter<'_>,
1278 ) -> Option<ColumnWriterImpl<'_, Self>> {
1279 match column_writer {
1280 ColumnWriter::$writer_ident(w) => Some(w),
1281 _ => None,
1282 }
1283 }
1284
1285 fn get_column_writer_ref<'a, 'b: 'a>(
1286 column_writer: &'a ColumnWriter<'b>,
1287 ) -> Option<&'a ColumnWriterImpl<'b, Self>> {
1288 match column_writer {
1289 ColumnWriter::$writer_ident(w) => Some(w),
1290 _ => None,
1291 }
1292 }
1293
1294 fn get_column_writer_mut<'a, 'b: 'a>(
1295 column_writer: &'a mut ColumnWriter<'b>,
1296 ) -> Option<&'a mut ColumnWriterImpl<'b, Self>> {
1297 match column_writer {
1298 ColumnWriter::$writer_ident(w) => Some(w),
1299 _ => None,
1300 }
1301 }
1302 }
1303 };
1304}
1305
1306make_type!(BoolType, BoolColumnReader, BoolColumnWriter, bool, 1);
1309make_type!(Int32Type, Int32ColumnReader, Int32ColumnWriter, i32, 4);
1310make_type!(Int64Type, Int64ColumnReader, Int64ColumnWriter, i64, 8);
1311make_type!(
1312 Int96Type,
1313 Int96ColumnReader,
1314 Int96ColumnWriter,
1315 Int96,
1316 mem::size_of::<Int96>()
1317);
1318make_type!(FloatType, FloatColumnReader, FloatColumnWriter, f32, 4);
1319make_type!(DoubleType, DoubleColumnReader, DoubleColumnWriter, f64, 8);
1320make_type!(
1321 ByteArrayType,
1322 ByteArrayColumnReader,
1323 ByteArrayColumnWriter,
1324 ByteArray,
1325 mem::size_of::<ByteArray>()
1326);
1327make_type!(
1328 FixedLenByteArrayType,
1329 FixedLenByteArrayColumnReader,
1330 FixedLenByteArrayColumnWriter,
1331 FixedLenByteArray,
1332 mem::size_of::<FixedLenByteArray>()
1333);
1334
1335impl AsRef<[u8]> for ByteArray {
1336 fn as_ref(&self) -> &[u8] {
1337 self.as_bytes()
1338 }
1339}
1340
1341impl AsRef<[u8]> for FixedLenByteArray {
1342 fn as_ref(&self) -> &[u8] {
1343 self.as_bytes()
1344 }
1345}
1346
1347macro_rules! ensure_phys_ty {
1349 ($($ty:pat_param)|+ , $($arg:tt)*) => {
1350 match T::get_physical_type() {
1351 $($ty => (),)*
1352 _ => panic!($($arg)*),
1353 };
1354 }
1355}
1356
1357#[cfg(test)]
1358mod tests {
1359 use super::*;
1360
1361 #[test]
1362 fn test_as_bytes() {
1363 let i96 = Int96::from(vec![1, 2, 3]);
1365 assert_eq!(i96.as_bytes(), &[1, 0, 0, 0, 2, 0, 0, 0, 3, 0, 0, 0]);
1366
1367 let byte_arr = ByteArray::from(vec![1, 2, 3]);
1369 assert_eq!(byte_arr.as_bytes(), &[1, 2, 3]);
1370
1371 let decimal = Decimal::from_i32(123, 5, 2);
1373 assert_eq!(decimal.as_bytes(), &[0, 0, 0, 123]);
1374 let decimal = Decimal::from_i64(123, 5, 2);
1375 assert_eq!(decimal.as_bytes(), &[0, 0, 0, 0, 0, 0, 0, 123]);
1376 let decimal = Decimal::from_bytes(ByteArray::from(vec![1, 2, 3]), 5, 2);
1377 assert_eq!(decimal.as_bytes(), &[1, 2, 3]);
1378 }
1379
1380 #[test]
1381 fn test_int96_from() {
1382 assert_eq!(
1383 Int96::from(vec![1, 12345, 1234567890]).data(),
1384 &[1, 12345, 1234567890]
1385 );
1386 }
1387
1388 #[test]
1389 fn test_byte_array_from() {
1390 assert_eq!(ByteArray::from(b"ABC".to_vec()).data(), b"ABC");
1391 assert_eq!(ByteArray::from("ABC").data(), b"ABC");
1392 assert_eq!(
1393 ByteArray::from(Bytes::from(vec![1u8, 2u8, 3u8, 4u8, 5u8])).data(),
1394 &[1u8, 2u8, 3u8, 4u8, 5u8]
1395 );
1396 let buf = vec![6u8, 7u8, 8u8, 9u8, 10u8];
1397 assert_eq!(ByteArray::from(buf).data(), &[6u8, 7u8, 8u8, 9u8, 10u8]);
1398 }
1399
1400 #[test]
1401 fn test_decimal_partial_eq() {
1402 assert_eq!(Decimal::default(), Decimal::from_i32(0, 0, 0));
1403 assert_eq!(Decimal::from_i32(222, 5, 2), Decimal::from_i32(222, 5, 2));
1404 assert_eq!(
1405 Decimal::from_bytes(ByteArray::from(vec![0, 0, 0, 3]), 5, 2),
1406 Decimal::from_i32(3, 5, 2)
1407 );
1408
1409 assert_ne!(Decimal::from_i32(222, 5, 2), Decimal::from_i32(111, 5, 2));
1410 assert_ne!(Decimal::from_i32(222, 5, 2), Decimal::from_i32(222, 6, 2));
1411 assert_ne!(Decimal::from_i32(222, 5, 2), Decimal::from_i32(222, 5, 3));
1412
1413 assert_ne!(Decimal::from_i64(222, 5, 2), Decimal::from_i32(222, 5, 2));
1414 }
1415
1416 #[test]
1417 fn test_byte_array_ord() {
1418 let byte_arr1 = ByteArray::from(vec![1, 2, 3]);
1419 let byte_arr11 = ByteArray::from(vec![1, 2, 3]);
1420 let byte_arr2 = ByteArray::from(vec![3, 4]);
1421 let byte_arr3 = ByteArray::from(vec![1, 2, 4]);
1422 let byte_arr4 = ByteArray::from(vec![]);
1423 let byte_arr5 = ByteArray::from(vec![2, 2, 3]);
1424
1425 assert!(byte_arr1 < byte_arr2);
1426 assert!(byte_arr3 > byte_arr1);
1427 assert!(byte_arr1 > byte_arr4);
1428 assert_eq!(byte_arr1, byte_arr11);
1429 assert!(byte_arr5 > byte_arr1);
1430 }
1431}