parquet/
parquet_thrift.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
18//! Structs used for encoding and decoding Parquet Thrift objects.
19//!
20//! These include:
21//! * [`ThriftCompactInputProtocol`]: Trait implemented by Thrift decoders.
22//!     * [`ThriftSliceInputProtocol`]: Thrift decoder that takes a slice of bytes as input.
23//!     * [`ThriftReadInputProtocol`]: Thrift decoder that takes a [`Read`] as input.
24//! * [`ReadThrift`]: Trait implemented by serializable objects.
25//! * [`ThriftCompactOutputProtocol`]: Thrift encoder.
26//! * [`WriteThrift`]: Trait implemented by serializable objects.
27//! * [`WriteThriftField`]: Trait implemented by serializable objects that are fields in Thrift structs.
28
29use std::{
30    cmp::Ordering,
31    io::{Read, Write},
32};
33
34use crate::{
35    errors::{ParquetError, Result},
36    write_thrift_field,
37};
38use std::io::Error;
39use std::num::TryFromIntError;
40use std::str::Utf8Error;
41
42#[derive(Debug)]
43pub(crate) enum ThriftProtocolError {
44    Eof,
45    IO(Error),
46    InvalidFieldType(u8),
47    InvalidElementType(u8),
48    FieldDeltaOverflow { field_delta: u8, last_field_id: i16 },
49    InvalidBoolean(u8),
50    IntegerOverflow,
51    Utf8Error,
52    SkipDepth(FieldType),
53    SkipUnsupportedType(FieldType),
54}
55
56impl From<ThriftProtocolError> for ParquetError {
57    #[inline(never)]
58    fn from(e: ThriftProtocolError) -> Self {
59        match e {
60            ThriftProtocolError::Eof => eof_err!("Unexpected EOF"),
61            ThriftProtocolError::IO(e) => e.into(),
62            ThriftProtocolError::InvalidFieldType(value) => match FieldType::try_from(value) {
63                Ok(fld_type) => general_err!("Unexpected struct field type {:?}", fld_type),
64                Err(_) => general_err!("Unexpected struct field type {}", value),
65            },
66            ThriftProtocolError::InvalidElementType(value) => {
67                general_err!("Unexpected list/set element type {}", value)
68            }
69            ThriftProtocolError::FieldDeltaOverflow {
70                field_delta,
71                last_field_id,
72            } => general_err!("cannot add {} to {}", field_delta, last_field_id),
73            ThriftProtocolError::InvalidBoolean(value) => {
74                general_err!("cannot convert {} into bool", value)
75            }
76            ThriftProtocolError::IntegerOverflow => {
77                general_err!("integer overflow decoding thrift value")
78            }
79            ThriftProtocolError::Utf8Error => general_err!("invalid utf8"),
80            ThriftProtocolError::SkipDepth(field_type) => {
81                general_err!("cannot parse past {:?}", field_type)
82            }
83            ThriftProtocolError::SkipUnsupportedType(field_type) => {
84                general_err!("cannot skip field type {:?}", field_type)
85            }
86        }
87    }
88}
89
90impl From<Utf8Error> for ThriftProtocolError {
91    fn from(_: Utf8Error) -> Self {
92        // ignore error payload to reduce the size of ThriftProtocolError
93        Self::Utf8Error
94    }
95}
96
97impl From<Error> for ThriftProtocolError {
98    fn from(e: Error) -> Self {
99        Self::IO(e)
100    }
101}
102
103impl From<TryFromIntError> for ThriftProtocolError {
104    fn from(_: TryFromIntError) -> Self {
105        // ignore error payload to reduce the size of ThriftProtocolError
106        Self::IntegerOverflow
107    }
108}
109
110pub type ThriftProtocolResult<T> = Result<T, ThriftProtocolError>;
111
112/// Wrapper for thrift `double` fields. This is used to provide
113/// an implementation of `Eq` for floats. This implementation
114/// uses IEEE 754 total order.
115#[derive(Debug, Clone, Copy, PartialEq)]
116pub struct OrderedF64(f64);
117
118impl From<f64> for OrderedF64 {
119    fn from(value: f64) -> Self {
120        Self(value)
121    }
122}
123
124impl From<OrderedF64> for f64 {
125    fn from(value: OrderedF64) -> Self {
126        value.0
127    }
128}
129
130impl Eq for OrderedF64 {} // Marker trait, requires PartialEq
131
132impl Ord for OrderedF64 {
133    fn cmp(&self, other: &Self) -> Ordering {
134        self.0.total_cmp(&other.0)
135    }
136}
137
138impl PartialOrd for OrderedF64 {
139    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
140        Some(self.cmp(other))
141    }
142}
143
144// Thrift compact protocol types for struct fields.
145#[derive(Clone, Copy, Debug, Eq, PartialEq)]
146pub(crate) enum FieldType {
147    Stop = 0,
148    BooleanTrue = 1,
149    BooleanFalse = 2,
150    Byte = 3,
151    I16 = 4,
152    I32 = 5,
153    I64 = 6,
154    Double = 7,
155    Binary = 8,
156    List = 9,
157    Set = 10,
158    Map = 11,
159    Struct = 12,
160    Uuid = 13,
161}
162
163impl TryFrom<u8> for FieldType {
164    type Error = ThriftProtocolError;
165    fn try_from(value: u8) -> ThriftProtocolResult<Self> {
166        match value {
167            0 => Ok(Self::Stop),
168            1 => Ok(Self::BooleanTrue),
169            2 => Ok(Self::BooleanFalse),
170            3 => Ok(Self::Byte),
171            4 => Ok(Self::I16),
172            5 => Ok(Self::I32),
173            6 => Ok(Self::I64),
174            7 => Ok(Self::Double),
175            8 => Ok(Self::Binary),
176            9 => Ok(Self::List),
177            10 => Ok(Self::Set),
178            11 => Ok(Self::Map),
179            12 => Ok(Self::Struct),
180            13 => Ok(Self::Uuid),
181            _ => Err(ThriftProtocolError::InvalidFieldType(value)),
182        }
183    }
184}
185
186impl From<ElementType> for FieldType {
187    fn from(value: ElementType) -> Self {
188        match value {
189            ElementType::Bool => Self::BooleanTrue,
190            ElementType::Byte => Self::Byte,
191            ElementType::I16 => Self::I16,
192            ElementType::I32 => Self::I32,
193            ElementType::I64 => Self::I64,
194            ElementType::Double => Self::Double,
195            ElementType::Binary => Self::Binary,
196            ElementType::List => Self::List,
197            ElementType::Set => Self::Set,
198            ElementType::Map => Self::Map,
199            ElementType::Struct => Self::Struct,
200            ElementType::Uuid => Self::Uuid,
201        }
202    }
203}
204
205// Thrift compact protocol types for list elements
206#[derive(Clone, Copy, Debug, Eq, PartialEq)]
207pub(crate) enum ElementType {
208    Bool = 2,
209    Byte = 3,
210    I16 = 4,
211    I32 = 5,
212    I64 = 6,
213    Double = 7,
214    Binary = 8,
215    List = 9,
216    Set = 10,
217    Map = 11,
218    Struct = 12,
219    Uuid = 13,
220}
221
222impl TryFrom<u8> for ElementType {
223    type Error = ThriftProtocolError;
224    fn try_from(value: u8) -> ThriftProtocolResult<Self> {
225        match value {
226            // For historical and compatibility reasons, a reader should be capable to deal with both cases.
227            // The only valid value in the original spec was 2, but due to an widespread implementation bug
228            // the defacto standard across large parts of the library became 1 instead.
229            // As a result, both values are now allowed.
230            // https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set
231            1 | 2 => Ok(Self::Bool),
232            3 => Ok(Self::Byte),
233            4 => Ok(Self::I16),
234            5 => Ok(Self::I32),
235            6 => Ok(Self::I64),
236            7 => Ok(Self::Double),
237            8 => Ok(Self::Binary),
238            9 => Ok(Self::List),
239            10 => Ok(Self::Set),
240            11 => Ok(Self::Map),
241            12 => Ok(Self::Struct),
242            13 => Ok(Self::Uuid),
243            _ => Err(ThriftProtocolError::InvalidElementType(value)),
244        }
245    }
246}
247
248/// Struct used to describe a [thrift struct] field during decoding.
249///
250/// [thrift struct]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct-encoding
251pub(crate) struct FieldIdentifier {
252    /// The type for the field.
253    pub(crate) field_type: FieldType,
254    /// The field's `id`. May be computed from delta or directly decoded.
255    pub(crate) id: i16,
256}
257
258impl FieldIdentifier {
259    pub(crate) fn bool_val(&self) -> ThriftProtocolResult<bool> {
260        match self.field_type {
261            FieldType::BooleanTrue => Ok(true),
262            FieldType::BooleanFalse => Ok(false),
263            _ => Err(ThriftProtocolError::InvalidFieldType(self.field_type as u8)),
264        }
265    }
266}
267
268/// Struct used to describe a [thrift list].
269///
270/// [thrift list]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set
271#[derive(Clone, Debug, Eq, PartialEq)]
272pub(crate) struct ListIdentifier {
273    /// The type for each element in the list.
274    pub(crate) element_type: ElementType,
275    /// Number of elements contained in the list.
276    pub(crate) size: i32,
277}
278
279/// Low-level object used to deserialize structs encoded with the Thrift [compact] protocol.
280///
281/// Implementation of this trait must provide the low-level functions `read_byte`, `read_bytes`,
282/// `skip_bytes`, and `read_double`. These primitives are used by the default functions provided
283/// here to perform deserialization.
284///
285/// [compact]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md
286pub(crate) trait ThriftCompactInputProtocol<'a> {
287    /// Read a single byte from the input.
288    fn read_byte(&mut self) -> ThriftProtocolResult<u8>;
289
290    /// Read a Thrift encoded [binary] from the input.
291    ///
292    /// [binary]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#binary-encoding
293    fn read_bytes(&mut self) -> ThriftProtocolResult<&'a [u8]>;
294
295    fn read_bytes_owned(&mut self) -> ThriftProtocolResult<Vec<u8>>;
296
297    /// Skip the next `n` bytes of input.
298    fn skip_bytes(&mut self, n: usize) -> ThriftProtocolResult<()>;
299
300    /// Read a ULEB128 encoded unsigned varint from the input.
301    fn read_vlq(&mut self) -> ThriftProtocolResult<u64> {
302        // try the happy path first
303        let byte = self.read_byte()?;
304        if byte & 0x80 == 0 {
305            return Ok(byte as u64);
306        }
307        let mut in_progress = (byte & 0x7f) as u64;
308        let mut shift = 7;
309        loop {
310            let byte = self.read_byte()?;
311            in_progress |= ((byte & 0x7F) as u64).wrapping_shl(shift);
312            if byte & 0x80 == 0 {
313                return Ok(in_progress);
314            }
315            shift += 7;
316        }
317    }
318
319    /// Read a zig-zag encoded signed varint from the input.
320    fn read_zig_zag(&mut self) -> ThriftProtocolResult<i64> {
321        let val = self.read_vlq()?;
322        Ok((val >> 1) as i64 ^ -((val & 1) as i64))
323    }
324
325    /// Read the [`ListIdentifier`] for a Thrift encoded list.
326    fn read_list_begin(&mut self) -> ThriftProtocolResult<ListIdentifier> {
327        let header = self.read_byte()?;
328        // some parquet writers will have an element_type of 0 for an empty list.
329        // account for that and return a bogus but valid element_type.
330        if header == 0 {
331            return Ok(ListIdentifier {
332                element_type: ElementType::Byte,
333                size: 0,
334            });
335        }
336        let element_type = ElementType::try_from(header & 0x0f)?;
337
338        let possible_element_count = (header & 0xF0) >> 4;
339        let element_count = if possible_element_count != 15 {
340            // high bits set high if count and type encoded separately
341            possible_element_count as i32
342        } else {
343            // The list size on the wire is an unsigned varint, but we represent
344            // it as `i32` (matching Java's `int` and the Thrift schema).
345            // A varint that decodes above `i32::MAX` is malformed input — reject
346            // it here at the protocol layer rather than letting the cast wrap
347            // into a negative size that downstream allocation code has to
348            // re-validate.
349            i32::try_from(self.read_vlq()?)?
350        };
351
352        Ok(ListIdentifier {
353            element_type,
354            size: element_count,
355        })
356    }
357
358    // Full field ids are uncommon.
359    // Not inlining this method reduces the code size of `read_field_begin`, which then ideally gets
360    // inlined everywhere.
361    #[cold]
362    fn read_full_field_id(&mut self) -> ThriftProtocolResult<i16> {
363        self.read_i16()
364    }
365
366    /// Read the [`FieldIdentifier`] for a field in a Thrift encoded struct.
367    fn read_field_begin(&mut self, last_field_id: i16) -> ThriftProtocolResult<FieldIdentifier> {
368        // we can read at least one byte, which is:
369        // - the type
370        // - the field delta and the type
371        let field_type = self.read_byte()?;
372        if field_type & 0xf == 0 {
373            return Ok(FieldIdentifier {
374                field_type: FieldType::Stop,
375                id: 0,
376            });
377        }
378
379        let field_delta = (field_type & 0xf0) >> 4;
380        let field_type = FieldType::try_from(field_type & 0xf)?;
381
382        let id = if field_delta != 0 {
383            last_field_id.checked_add(field_delta as i16).ok_or(
384                ThriftProtocolError::FieldDeltaOverflow {
385                    field_delta,
386                    last_field_id,
387                },
388            )?
389        } else {
390            self.read_full_field_id()?
391        };
392
393        Ok(FieldIdentifier { field_type, id })
394    }
395
396    /// This is a specialized version of [`Self::read_field_begin`], solely for use in parsing
397    /// simple structs. This function assumes that the delta field will always be less than 0xf,
398    /// fields will be in order, and no boolean fields will be read.
399    /// This also skips validation of the field type.
400    ///
401    /// Returns a tuple of `(field_type, field_delta)`.
402    fn read_field_header(&mut self) -> ThriftProtocolResult<(u8, u8)> {
403        let field_type = self.read_byte()?;
404        let field_delta = (field_type & 0xf0) >> 4;
405        let field_type = field_type & 0xf;
406        Ok((field_type, field_delta))
407    }
408
409    /// Read a boolean list element. This should not be used for struct fields. For the latter,
410    /// use the [`FieldIdentifier::bool_val`] field.
411    fn read_bool(&mut self) -> ThriftProtocolResult<bool> {
412        let b = self.read_byte()?;
413        // Previous versions of the thrift specification said to use 0 and 1 inside collections,
414        // but that differed from existing implementations.
415        // The specification was updated in https://github.com/apache/thrift/commit/2c29c5665bc442e703480bb0ee60fe925ffe02e8.
416        // At least the go implementation seems to have followed the previously documented values.
417        match b {
418            0x01 => Ok(true),
419            0x00 | 0x02 => Ok(false),
420            _ => Err(ThriftProtocolError::InvalidBoolean(b)),
421        }
422    }
423
424    /// Read a Thrift [binary] as a UTF-8 encoded string.
425    ///
426    /// [binary]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#binary-encoding
427    fn read_string(&mut self) -> ThriftProtocolResult<&'a str> {
428        let slice = self.read_bytes()?;
429        Ok(std::str::from_utf8(slice)?)
430    }
431
432    /// Read an `i8`.
433    fn read_i8(&mut self) -> ThriftProtocolResult<i8> {
434        Ok(self.read_byte()? as _)
435    }
436
437    /// Read an `i16`.
438    fn read_i16(&mut self) -> ThriftProtocolResult<i16> {
439        Ok(self.read_zig_zag()? as _)
440    }
441
442    /// Read an `i32`.
443    fn read_i32(&mut self) -> ThriftProtocolResult<i32> {
444        Ok(self.read_zig_zag()? as _)
445    }
446
447    /// Read an `i64`.
448    fn read_i64(&mut self) -> ThriftProtocolResult<i64> {
449        self.read_zig_zag()
450    }
451
452    /// Read a Thrift `double` as `f64`.
453    fn read_double(&mut self) -> ThriftProtocolResult<f64>;
454
455    /// Skip a ULEB128 encoded varint.
456    fn skip_vlq(&mut self) -> ThriftProtocolResult<()> {
457        loop {
458            let byte = self.read_byte()?;
459            if byte & 0x80 == 0 {
460                return Ok(());
461            }
462        }
463    }
464
465    /// Skip a thrift [binary].
466    ///
467    /// [binary]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#binary-encoding
468    fn skip_binary(&mut self) -> ThriftProtocolResult<()> {
469        let len = self.read_vlq()? as usize;
470        self.skip_bytes(len)
471    }
472
473    /// Skip a field with type `field_type` recursively until the default
474    /// maximum skip depth (currently 64) is reached.
475    fn skip(&mut self, field_type: FieldType) -> ThriftProtocolResult<()> {
476        const DEFAULT_SKIP_DEPTH: i8 = 64;
477        self.skip_till_depth(field_type, DEFAULT_SKIP_DEPTH)
478    }
479
480    /// Empty structs in unions consist of a single byte of 0 for the field stop record.
481    /// This skips that byte without encuring the cost of processing the [`FieldIdentifier`].
482    /// Will return an error if the struct is not actually empty.
483    fn skip_empty_struct(&mut self) -> Result<()> {
484        let b = self.read_byte()?;
485        if b != 0 {
486            Err(general_err!("Empty struct has fields"))
487        } else {
488            Ok(())
489        }
490    }
491
492    /// Skip a field with type `field_type` recursively up to `depth` levels.
493    fn skip_till_depth(&mut self, field_type: FieldType, depth: i8) -> ThriftProtocolResult<()> {
494        if depth == 0 {
495            return Err(ThriftProtocolError::SkipDepth(field_type));
496        }
497
498        match field_type {
499            // boolean field has no data
500            FieldType::BooleanFalse | FieldType::BooleanTrue => Ok(()),
501            FieldType::Byte => self.read_i8().map(|_| ()),
502            FieldType::I16 => self.skip_vlq(),
503            FieldType::I32 => self.skip_vlq(),
504            FieldType::I64 => self.skip_vlq(),
505            FieldType::Double => self.skip_bytes(8),
506            FieldType::Binary => self.skip_binary(),
507            // see https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct
508            FieldType::Struct => {
509                loop {
510                    // we don't need field id for skipping, so always pass 0 for last id
511                    let field_ident = self.read_field_begin(0)?;
512                    if field_ident.field_type == FieldType::Stop {
513                        break;
514                    }
515                    self.skip_till_depth(field_ident.field_type, depth - 1)?;
516                }
517                Ok(())
518            }
519            // lists and sets are encoded the same
520            // see https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set
521            FieldType::List | FieldType::Set => {
522                let list_ident = self.read_list_begin()?;
523                let element_type = FieldType::from(list_ident.element_type);
524                for _ in 0..list_ident.size {
525                    self.skip_till_depth(element_type, depth - 1)?;
526                }
527                Ok(())
528            }
529            // see https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#map
530            FieldType::Map => {
531                let size = i32::try_from(self.read_vlq()?)?;
532                if size > 0 {
533                    let kv = self.read_byte()?;
534                    let key_type = FieldType::from(ElementType::try_from(kv >> 4)?);
535                    let val_type = FieldType::from(ElementType::try_from(kv & 0xf)?);
536                    for _ in 0..size {
537                        self.skip_till_depth(key_type, depth - 1)?;
538                        self.skip_till_depth(val_type, depth - 1)?;
539                    }
540                }
541                Ok(())
542            }
543            // see https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#universal-unique-identifier-encoding
544            FieldType::Uuid => self.skip_bytes(16),
545            _ => Err(ThriftProtocolError::SkipUnsupportedType(field_type)),
546        }
547    }
548}
549
550/// A high performance Thrift reader that reads from a slice of bytes.
551pub(crate) struct ThriftSliceInputProtocol<'a> {
552    buf: &'a [u8],
553}
554
555impl<'a> ThriftSliceInputProtocol<'a> {
556    /// Create a new `ThriftSliceInputProtocol` using the bytes in `buf`.
557    pub fn new(buf: &'a [u8]) -> Self {
558        Self { buf }
559    }
560
561    /// Return the current buffer as a slice.
562    pub fn as_slice(&self) -> &'a [u8] {
563        self.buf
564    }
565}
566
567impl<'b, 'a: 'b> ThriftCompactInputProtocol<'b> for ThriftSliceInputProtocol<'a> {
568    #[inline]
569    fn read_byte(&mut self) -> ThriftProtocolResult<u8> {
570        let ret = *self.buf.first().ok_or(ThriftProtocolError::Eof)?;
571        self.buf = &self.buf[1..];
572        Ok(ret)
573    }
574
575    fn read_bytes(&mut self) -> ThriftProtocolResult<&'b [u8]> {
576        let len = self.read_vlq()? as usize;
577        let ret = self.buf.get(..len).ok_or(ThriftProtocolError::Eof)?;
578        self.buf = &self.buf[len..];
579        Ok(ret)
580    }
581
582    fn read_bytes_owned(&mut self) -> ThriftProtocolResult<Vec<u8>> {
583        Ok(self.read_bytes()?.to_vec())
584    }
585
586    #[inline]
587    fn skip_bytes(&mut self, n: usize) -> ThriftProtocolResult<()> {
588        self.buf.get(..n).ok_or(ThriftProtocolError::Eof)?;
589        self.buf = &self.buf[n..];
590        Ok(())
591    }
592
593    fn read_double(&mut self) -> ThriftProtocolResult<f64> {
594        let slice = self.buf.get(..8).ok_or(ThriftProtocolError::Eof)?;
595        self.buf = &self.buf[8..];
596        match slice.try_into() {
597            Ok(slice) => Ok(f64::from_le_bytes(slice)),
598            Err(_) => unreachable!(),
599        }
600    }
601}
602
603/// A Thrift input protocol that wraps a [`Read`] object.
604///
605/// Note that this is only intended for use in reading Parquet page headers. This will panic
606/// if Thrift `binary` data is encountered because a slice of that data cannot be returned.
607pub(crate) struct ThriftReadInputProtocol<R: Read> {
608    reader: R,
609}
610
611impl<R: Read> ThriftReadInputProtocol<R> {
612    pub(crate) fn new(reader: R) -> Self {
613        Self { reader }
614    }
615}
616
617impl<'a, R: Read> ThriftCompactInputProtocol<'a> for ThriftReadInputProtocol<R> {
618    #[inline]
619    fn read_byte(&mut self) -> ThriftProtocolResult<u8> {
620        let mut buf = [0_u8; 1];
621        self.reader.read_exact(&mut buf)?;
622        Ok(buf[0])
623    }
624
625    fn read_bytes(&mut self) -> ThriftProtocolResult<&'a [u8]> {
626        unimplemented!()
627    }
628
629    fn read_bytes_owned(&mut self) -> ThriftProtocolResult<Vec<u8>> {
630        let len = self.read_vlq()? as usize;
631        let mut v = Vec::with_capacity(len);
632        std::io::copy(&mut self.reader.by_ref().take(len as u64), &mut v)?;
633        Ok(v)
634    }
635
636    fn skip_bytes(&mut self, n: usize) -> ThriftProtocolResult<()> {
637        std::io::copy(
638            &mut self.reader.by_ref().take(n as u64),
639            &mut std::io::sink(),
640        )?;
641        Ok(())
642    }
643
644    fn read_double(&mut self) -> ThriftProtocolResult<f64> {
645        let mut buf = [0_u8; 8];
646        self.reader.read_exact(&mut buf)?;
647        Ok(f64::from_le_bytes(buf))
648    }
649}
650
651/// Trait implemented for objects that can be deserialized from a Thrift input stream.
652/// Implementations are provided for Thrift primitive types.
653pub(crate) trait ReadThrift<'a, R: ThriftCompactInputProtocol<'a>> {
654    /// Read an object of type `Self` from the input protocol object.
655    fn read_thrift(prot: &mut R) -> Result<Self>
656    where
657        Self: Sized;
658}
659
660impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for bool {
661    fn read_thrift(prot: &mut R) -> Result<Self> {
662        Ok(prot.read_bool()?)
663    }
664}
665
666impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for i8 {
667    fn read_thrift(prot: &mut R) -> Result<Self> {
668        Ok(prot.read_i8()?)
669    }
670}
671
672impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for i16 {
673    fn read_thrift(prot: &mut R) -> Result<Self> {
674        Ok(prot.read_i16()?)
675    }
676}
677
678impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for i32 {
679    fn read_thrift(prot: &mut R) -> Result<Self> {
680        Ok(prot.read_i32()?)
681    }
682}
683
684impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for i64 {
685    fn read_thrift(prot: &mut R) -> Result<Self> {
686        Ok(prot.read_i64()?)
687    }
688}
689
690impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for OrderedF64 {
691    fn read_thrift(prot: &mut R) -> Result<Self> {
692        Ok(OrderedF64(prot.read_double()?))
693    }
694}
695
696impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for &'a str {
697    fn read_thrift(prot: &mut R) -> Result<Self> {
698        Ok(prot.read_string()?)
699    }
700}
701
702impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for String {
703    fn read_thrift(prot: &mut R) -> Result<Self> {
704        Ok(String::from_utf8(prot.read_bytes_owned()?)?)
705    }
706}
707
708impl<'a, R: ThriftCompactInputProtocol<'a>> ReadThrift<'a, R> for &'a [u8] {
709    fn read_thrift(prot: &mut R) -> Result<Self> {
710        Ok(prot.read_bytes()?)
711    }
712}
713
714/// Read a Thrift encoded [list] from the input protocol object.
715///
716/// [list]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#list-and-set
717pub(crate) fn read_thrift_vec<'a, T, R>(prot: &mut R) -> Result<Vec<T>>
718where
719    R: ThriftCompactInputProtocol<'a>,
720    T: ReadThrift<'a, R> + WriteThrift,
721{
722    let list_ident = prot.read_list_begin()?;
723    validate_list_type(T::ELEMENT_TYPE, &list_ident)?;
724    let mut res = Vec::with_capacity(list_ident.size as usize);
725    for _ in 0..list_ident.size {
726        let val = T::read_thrift(prot)?;
727        res.push(val);
728    }
729    Ok(res)
730}
731
732pub(crate) fn validate_list_type(expected: ElementType, got: &ListIdentifier) -> Result<()> {
733    if got.element_type != expected {
734        return Err(general_err!(
735            "Expected list element type of {:?} but got {:?}",
736            expected,
737            got.element_type
738        ));
739    }
740    Ok(())
741}
742
743/////////////////////////
744// thrift compact output
745
746/// Low-level object used to serialize structs to the Thrift [compact output] protocol.
747///
748/// This struct serves as a wrapper around a [`Write`] object, to which thrift encoded data
749/// will written. The implementation provides functions to write Thrift primitive types, as well
750/// as functions used in the encoding of lists and structs. This should rarely be used directly,
751/// but is instead intended for use by implementers of [`WriteThrift`] and [`WriteThriftField`].
752///
753/// [compact output]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md
754pub(crate) struct ThriftCompactOutputProtocol<W: Write> {
755    writer: W,
756    write_path_in_schema: bool,
757    write_rg_ordinal: bool,
758}
759
760impl<W: Write> ThriftCompactOutputProtocol<W> {
761    /// Create a new `ThriftCompactOutputProtocol` wrapping the byte sink `writer`.
762    pub(crate) fn new(writer: W) -> Self {
763        Self {
764            writer,
765            write_path_in_schema: true,
766            write_rg_ordinal: true,
767        }
768    }
769
770    // TODO(ets): at some point there should probably be a properties object
771    // to control aspects of thrift output. But since this is the only option to date
772    // I'm choosing a simpler API.
773    /// Control the writing of the `path_in_schema` element of the `ColumnMetaData`
774    pub(crate) fn set_write_path_in_schema(&mut self, val: bool) {
775        self.write_path_in_schema = val;
776    }
777
778    /// Indicate whether or not to emit `path_in_schema`.
779    pub(crate) fn write_path_in_schema(&self) -> bool {
780        self.write_path_in_schema
781    }
782
783    /// Control the writing of the `ordinal` element of the `RowGroup` struct.
784    ///
785    /// The Thrift `ordinal` field on the `RowGroup` struct is `i16`, but the
786    /// Thrift compact protocol allows for up to 2^31 elements in a list. If
787    /// more than 2^15 row groups are to be written, this can be set to `false`
788    /// to prevent writing the ordinal for some row groups but not others.
789    pub(crate) fn set_write_row_group_ordinal(&mut self, val: bool) {
790        self.write_rg_ordinal = val;
791    }
792
793    /// Indicate whether or not to emit `ordinal`.
794    pub(crate) fn write_row_group_ordinal(&self) -> bool {
795        self.write_rg_ordinal
796    }
797
798    /// Write a single byte to the output stream.
799    fn write_byte(&mut self, b: u8) -> Result<()> {
800        self.writer.write_all(&[b])?;
801        Ok(())
802    }
803
804    /// Write the given `u64` as a ULEB128 encoded varint.
805    fn write_vlq(&mut self, val: u64) -> Result<()> {
806        let mut v = val;
807        while v > 0x7f {
808            self.write_byte(v as u8 | 0x80)?;
809            v >>= 7;
810        }
811        self.write_byte(v as u8)
812    }
813
814    /// Write the given `i64` as a zig-zag encoded varint.
815    fn write_zig_zag(&mut self, val: i64) -> Result<()> {
816        let s = (val < 0) as i64;
817        self.write_vlq((((val ^ -s) << 1) + s) as u64)
818    }
819
820    /// Used to mark the start of a Thrift struct field of type `field_type`. `last_field_id`
821    /// is used to compute a delta to the given `field_id` per the compact protocol [spec].
822    ///
823    /// [spec]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct-encoding
824    pub(crate) fn write_field_begin(
825        &mut self,
826        field_type: FieldType,
827        field_id: i16,
828        last_field_id: i16,
829    ) -> Result<()> {
830        let delta = field_id.wrapping_sub(last_field_id);
831        if delta > 0 && delta <= 0xf {
832            self.write_byte(((delta as u8) << 4) | field_type as u8)
833        } else {
834            self.write_byte(field_type as u8)?;
835            self.write_i16(field_id)
836        }
837    }
838
839    /// Used to indicate the start of a list of `element_type` elements.
840    pub(crate) fn write_list_begin(&mut self, element_type: ElementType, len: usize) -> Result<()> {
841        if len < 15 {
842            self.write_byte(((len as u8) << 4) | element_type as u8)
843        } else {
844            self.write_byte(0xf0u8 | element_type as u8)?;
845            self.write_vlq(len as _)
846        }
847    }
848
849    /// Used to mark the end of a struct. This must be called after all fields of the struct have
850    /// been written.
851    pub(crate) fn write_struct_end(&mut self) -> Result<()> {
852        self.write_byte(0)
853    }
854
855    /// Serialize a slice of `u8`s. This will encode a length, and then write the bytes without
856    /// further encoding.
857    pub(crate) fn write_bytes(&mut self, val: &[u8]) -> Result<()> {
858        self.write_vlq(val.len() as u64)?;
859        self.writer.write_all(val)?;
860        Ok(())
861    }
862
863    /// Short-cut method used to encode structs that have no fields (often used in Thrift unions).
864    /// This simply encodes the field id and then immediately writes the end-of-struct marker.
865    pub(crate) fn write_empty_struct(&mut self, field_id: i16, last_field_id: i16) -> Result<i16> {
866        self.write_field_begin(FieldType::Struct, field_id, last_field_id)?;
867        self.write_struct_end()?;
868        Ok(last_field_id)
869    }
870
871    /// Write a boolean value.
872    pub(crate) fn write_bool(&mut self, val: bool) -> Result<()> {
873        match val {
874            true => self.write_byte(1),
875            false => self.write_byte(2),
876        }
877    }
878
879    /// Write a zig-zag encoded `i8` value.
880    pub(crate) fn write_i8(&mut self, val: i8) -> Result<()> {
881        self.write_byte(val as u8)
882    }
883
884    /// Write a zig-zag encoded `i16` value.
885    pub(crate) fn write_i16(&mut self, val: i16) -> Result<()> {
886        self.write_zig_zag(val as _)
887    }
888
889    /// Write a zig-zag encoded `i32` value.
890    pub(crate) fn write_i32(&mut self, val: i32) -> Result<()> {
891        self.write_zig_zag(val as _)
892    }
893
894    /// Write a zig-zag encoded `i64` value.
895    pub(crate) fn write_i64(&mut self, val: i64) -> Result<()> {
896        self.write_zig_zag(val)
897    }
898
899    /// Write a double value.
900    pub(crate) fn write_double(&mut self, val: f64) -> Result<()> {
901        self.writer.write_all(&val.to_le_bytes())?;
902        Ok(())
903    }
904}
905
906/// Trait implemented by objects that are to be serialized to a Thrift [compact output] protocol
907/// stream. Implementations are also provided for primitive Thrift types.
908///
909/// [compact output]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md
910pub(crate) trait WriteThrift {
911    /// The [`ElementType`] to use when a list of this object is written.
912    const ELEMENT_TYPE: ElementType;
913
914    /// Serialize this object to the given `writer`.
915    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()>;
916}
917
918/// Implementation for a vector of thrift serializable objects that implement [`WriteThrift`].
919/// This will write the necessary list header and then serialize the elements one-at-a-time.
920impl<T> WriteThrift for Vec<T>
921where
922    T: WriteThrift,
923{
924    const ELEMENT_TYPE: ElementType = ElementType::List;
925
926    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
927        writer.write_list_begin(T::ELEMENT_TYPE, self.len())?;
928        for item in self {
929            item.write_thrift(writer)?;
930        }
931        Ok(())
932    }
933}
934
935impl WriteThrift for bool {
936    const ELEMENT_TYPE: ElementType = ElementType::Bool;
937
938    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
939        writer.write_bool(*self)
940    }
941}
942
943impl WriteThrift for i8 {
944    const ELEMENT_TYPE: ElementType = ElementType::Byte;
945
946    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
947        writer.write_i8(*self)
948    }
949}
950
951impl WriteThrift for i16 {
952    const ELEMENT_TYPE: ElementType = ElementType::I16;
953
954    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
955        writer.write_i16(*self)
956    }
957}
958
959impl WriteThrift for i32 {
960    const ELEMENT_TYPE: ElementType = ElementType::I32;
961
962    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
963        writer.write_i32(*self)
964    }
965}
966
967impl WriteThrift for i64 {
968    const ELEMENT_TYPE: ElementType = ElementType::I64;
969
970    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
971        writer.write_i64(*self)
972    }
973}
974
975impl WriteThrift for OrderedF64 {
976    const ELEMENT_TYPE: ElementType = ElementType::Double;
977
978    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
979        writer.write_double(self.0)
980    }
981}
982
983impl WriteThrift for f64 {
984    const ELEMENT_TYPE: ElementType = ElementType::Double;
985
986    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
987        writer.write_double(*self)
988    }
989}
990
991impl WriteThrift for &[u8] {
992    const ELEMENT_TYPE: ElementType = ElementType::Binary;
993
994    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
995        writer.write_bytes(self)
996    }
997}
998
999impl WriteThrift for &str {
1000    const ELEMENT_TYPE: ElementType = ElementType::Binary;
1001
1002    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
1003        writer.write_bytes(self.as_bytes())
1004    }
1005}
1006
1007impl WriteThrift for String {
1008    const ELEMENT_TYPE: ElementType = ElementType::Binary;
1009
1010    fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
1011        writer.write_bytes(self.as_bytes())
1012    }
1013}
1014
1015/// Trait implemented by objects that are fields of Thrift structs.
1016///
1017/// For example, given the Thrift struct definition
1018/// ```ignore
1019/// struct MyStruct {
1020///   1: required i32 field1
1021///   2: optional bool field2
1022///   3: optional OtherStruct field3
1023/// }
1024/// ```
1025///
1026/// which becomes in Rust
1027/// ```ignore
1028/// # struct OtherStruct {}
1029/// struct MyStruct {
1030///   field1: i32,
1031///   field2: Option<bool>,
1032///   field3: Option<OtherStruct>,
1033/// }
1034/// ```
1035/// the impl of `WriteThrift` for `MyStruct` will use the `WriteThriftField` impls for `i32`,
1036/// `bool`, and `OtherStruct`.
1037///
1038/// ```ignore
1039/// impl WriteThrift for MyStruct {
1040///   fn write_thrift<W: Write>(&self, writer: &mut ThriftCompactOutputProtocol<W>) -> Result<()> {
1041///     let mut last_field_id = 0i16;
1042///     last_field_id = self.field1.write_thrift_field(writer, 1, last_field_id)?;
1043///     if self.field2.is_some() {
1044///       // if field2 is `None` then this assignment won't happen and last_field_id will remain
1045///       // `1` when writing `field3`
1046///       last_field_id = self.field2.write_thrift_field(writer, 2, last_field_id)?;
1047///     }
1048///     if self.field3.is_some() {
1049///       // no need to assign last_field_id since this is the final field.
1050///       self.field3.write_thrift_field(writer, 3, last_field_id)?;
1051///     }
1052///     writer.write_struct_end()
1053///   }
1054/// }
1055/// ```
1056///
1057pub(crate) trait WriteThriftField {
1058    /// Used to write struct fields (which may be primitive or IDL defined types). This will
1059    /// write the field marker for the given `field_id`, using `last_field_id` to compute the
1060    /// field delta used by the Thrift [compact protocol]. On success this will return `field_id`
1061    /// to be used in chaining.
1062    ///
1063    /// [compact protocol]: https://github.com/apache/thrift/blob/master/doc/specs/thrift-compact-protocol.md#struct-encoding
1064    fn write_thrift_field<W: Write>(
1065        &self,
1066        writer: &mut ThriftCompactOutputProtocol<W>,
1067        field_id: i16,
1068        last_field_id: i16,
1069    ) -> Result<i16>;
1070}
1071
1072// bool struct fields are written differently to bool values
1073impl WriteThriftField for bool {
1074    fn write_thrift_field<W: Write>(
1075        &self,
1076        writer: &mut ThriftCompactOutputProtocol<W>,
1077        field_id: i16,
1078        last_field_id: i16,
1079    ) -> Result<i16> {
1080        // boolean only writes the field header
1081        match *self {
1082            true => writer.write_field_begin(FieldType::BooleanTrue, field_id, last_field_id)?,
1083            false => writer.write_field_begin(FieldType::BooleanFalse, field_id, last_field_id)?,
1084        }
1085        Ok(field_id)
1086    }
1087}
1088
1089write_thrift_field!(i8, FieldType::Byte);
1090write_thrift_field!(i16, FieldType::I16);
1091write_thrift_field!(i32, FieldType::I32);
1092write_thrift_field!(i64, FieldType::I64);
1093write_thrift_field!(OrderedF64, FieldType::Double);
1094write_thrift_field!(f64, FieldType::Double);
1095write_thrift_field!(String, FieldType::Binary);
1096
1097impl WriteThriftField for &[u8] {
1098    fn write_thrift_field<W: Write>(
1099        &self,
1100        writer: &mut ThriftCompactOutputProtocol<W>,
1101        field_id: i16,
1102        last_field_id: i16,
1103    ) -> Result<i16> {
1104        writer.write_field_begin(FieldType::Binary, field_id, last_field_id)?;
1105        writer.write_bytes(self)?;
1106        Ok(field_id)
1107    }
1108}
1109
1110impl WriteThriftField for &str {
1111    fn write_thrift_field<W: Write>(
1112        &self,
1113        writer: &mut ThriftCompactOutputProtocol<W>,
1114        field_id: i16,
1115        last_field_id: i16,
1116    ) -> Result<i16> {
1117        writer.write_field_begin(FieldType::Binary, field_id, last_field_id)?;
1118        writer.write_bytes(self.as_bytes())?;
1119        Ok(field_id)
1120    }
1121}
1122
1123impl<T> WriteThriftField for Vec<T>
1124where
1125    T: WriteThrift,
1126{
1127    fn write_thrift_field<W: Write>(
1128        &self,
1129        writer: &mut ThriftCompactOutputProtocol<W>,
1130        field_id: i16,
1131        last_field_id: i16,
1132    ) -> Result<i16> {
1133        writer.write_field_begin(FieldType::List, field_id, last_field_id)?;
1134        self.write_thrift(writer)?;
1135        Ok(field_id)
1136    }
1137}
1138
1139#[cfg(test)]
1140pub(crate) mod tests {
1141    use crate::basic::{TimeUnit, Type};
1142
1143    use super::*;
1144    use std::fmt::Debug;
1145
1146    pub(crate) fn test_roundtrip<T>(val: T)
1147    where
1148        T: for<'a> ReadThrift<'a, ThriftSliceInputProtocol<'a>> + WriteThrift + PartialEq + Debug,
1149    {
1150        let mut buf = Vec::<u8>::new();
1151        {
1152            let mut writer = ThriftCompactOutputProtocol::new(&mut buf);
1153            val.write_thrift(&mut writer).unwrap();
1154        }
1155
1156        let mut prot = ThriftSliceInputProtocol::new(&buf);
1157        let read_val = T::read_thrift(&mut prot).unwrap();
1158        assert_eq!(val, read_val);
1159    }
1160
1161    #[test]
1162    fn test_enum_roundtrip() {
1163        test_roundtrip(Type::BOOLEAN);
1164        test_roundtrip(Type::INT32);
1165        test_roundtrip(Type::INT64);
1166        test_roundtrip(Type::INT96);
1167        test_roundtrip(Type::FLOAT);
1168        test_roundtrip(Type::DOUBLE);
1169        test_roundtrip(Type::BYTE_ARRAY);
1170        test_roundtrip(Type::FIXED_LEN_BYTE_ARRAY);
1171    }
1172
1173    #[test]
1174    fn test_union_all_empty_roundtrip() {
1175        test_roundtrip(TimeUnit::MILLIS);
1176        test_roundtrip(TimeUnit::MICROS);
1177        test_roundtrip(TimeUnit::NANOS);
1178    }
1179
1180    #[test]
1181    fn test_decode_empty_list() {
1182        let data = vec![0u8; 1];
1183        let mut prot = ThriftSliceInputProtocol::new(&data);
1184        let header = prot.read_list_begin().expect("error reading list header");
1185        assert_eq!(header.size, 0);
1186        assert_eq!(header.element_type, ElementType::Byte);
1187    }
1188
1189    /// A Thrift list header whose `size` varint decodes above `i32::MAX`
1190    /// must be rejected at the protocol layer rather than wrapping into a
1191    /// negative `i32` and being smuggled into downstream allocation code.
1192    #[test]
1193    fn test_read_list_begin_size_above_i32_max_returns_err() {
1194        // List header: element_type=8 (Binary), 0xF=follow-up varint.
1195        // Varint 80 80 80 80 08 decodes to 0x8000_0000 = i32::MAX + 1.
1196        let mut data: Vec<u8> = vec![0xF8];
1197        data.extend_from_slice(&[0x80, 0x80, 0x80, 0x80, 0x08]);
1198        let mut prot = ThriftSliceInputProtocol::new(&data);
1199        let result = prot.read_list_begin();
1200        assert!(result.is_err(), "expected error, got {result:?}");
1201    }
1202
1203    #[test]
1204    fn test_read_list_wrong_type() {
1205        // list header: 4 elements of `Boolean`
1206        let data = [0x42, 0x01];
1207        let mut prot = ThriftSliceInputProtocol::new(&data);
1208        // try to read as list<i32>
1209        let result = read_thrift_vec::<i32, ThriftSliceInputProtocol>(&mut prot);
1210        println!("{result:?}");
1211        assert!(
1212            result
1213                .unwrap_err()
1214                .to_string()
1215                .contains("Expected list element type of I32 but got Bool")
1216        );
1217    }
1218}
parquet/parquet_thrift.rs

parquet/
parquet_thrift.rs
