Add option to skip validation when reading IPC streams/files

Author: alamb

arrow-ipc/src/reader.rs

@@ -40,6 +40,7 @@ use arrow_data::ArrayData;
 use arrow_schema::*;
 
 use crate::compression::CompressionCodec;
+use crate::reader::private::UnsafeFlag;
 use crate::{Block, FieldNode, Message, MetadataVersion, CONTINUATION_MARKER};
 use DataType::*;
 
@@ -65,6 +66,41 @@ fn read_buffer(
         (false, Some(decompressor)) => decompressor.decompress_to_buffer(&buf_data),
     }
 }
+
+mod private {
+    /// A boolean flag that cannot be mutated outside of unsafe code.
+    ///
+    /// Defaults to a value of false.
+    ///
+    /// This structure is used to enforce safety in the various readers
+    #[derive(Debug, Clone, Copy)]
+    pub struct UnsafeFlag(bool);
+
+    impl Default for UnsafeFlag {
+        fn default() -> Self {
+            Self::new()
+        }
+    }
+
+    impl UnsafeFlag {
+        /// Creates a new `UnsafeFlag` with the value set to `false`
+        #[inline]
+        pub const fn new() -> Self {
+            Self(false)
+        }
+
+        #[inline]
+        pub unsafe fn set(&mut self, val: bool) {
+            self.0 = val;
+        }
+
+        #[inline]
+        pub fn get(&self) -> bool {
+            self.0
+        }
+    }
+}
+
 impl RecordBatchDecoder<'_> {
     /// Coordinates reading arrays based on data types.
     ///
@@ -376,6 +412,18 @@ struct RecordBatchDecoder<'a> {
     /// Are buffers required to already be aligned? See
     /// [`RecordBatchDecoder::with_require_alignment`] for details
     require_alignment: bool,
+    /// Should validation be skipped when reading data?
+    ///
+    /// Defaults to false.
+    ///
+    /// If true [`ArrayData::validate`] is not called after reading
+    ///
+    /// # Safety
+    ///
+    /// This flag can only be set to true using `unsafe` APIs. However, once true
+    /// subsequent calls to `build()` may result in undefined behavior if the data
+    /// is not valid.
+    skip_validation: UnsafeFlag,
 }
 
 impl<'a> RecordBatchDecoder<'a> {
@@ -410,6 +458,7 @@ impl<'a> RecordBatchDecoder<'a> {
             buffers: buffers.iter(),
             projection: None,
             require_alignment: false,
+            skip_validation: UnsafeFlag::new(),
         })
     }
 
@@ -432,6 +481,21 @@ impl<'a> RecordBatchDecoder<'a> {
         self
     }
 
+    /// Set skip_validation (default: false)
+    ///
+    /// Note this is a pub(crate) API and can not be used outside of this crate
+    ///
+    /// If true, validation is skipped.
+    ///
+    /// # Safety
+    ///
+    /// Relies on `UnsafeFlag` to enforce safety -- can only be enabled via
+    /// unsafe APIs.
+    pub(crate) fn with_skip_validation(mut self, skip_validation: UnsafeFlag) -> Self {
+        self.skip_validation = skip_validation;
+        self
+    }
+
     /// Read the record batch, consuming the reader
     fn read_record_batch(mut self) -> Result<RecordBatch, ArrowError> {
         let mut variadic_counts: VecDeque<i64> = self
@@ -601,7 +665,16 @@ pub fn read_dictionary(
     dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
     metadata: &MetadataVersion,
 ) -> Result<(), ArrowError> {
-    read_dictionary_impl(buf, batch, schema, dictionaries_by_id, metadata, false)
+    let skip_validation = UnsafeFlag::new(); // do not skip valididation
+    read_dictionary_impl(
+        buf,
+        batch,
+        schema,
+        dictionaries_by_id,
+        metadata,
+        false,
+        skip_validation,
+    )
 }
 
 fn read_dictionary_impl(
@@ -611,6 +684,7 @@ fn read_dictionary_impl(
     dictionaries_by_id: &mut HashMap<i64, ArrayRef>,
     metadata: &MetadataVersion,
     require_alignment: bool,
+    skip_validation: UnsafeFlag,
 ) -> Result<(), ArrowError> {
     if batch.isDelta() {
         return Err(ArrowError::InvalidArgumentError(
@@ -642,6 +716,7 @@ fn read_dictionary_impl(
                 metadata,
             )?
             .with_require_alignment(require_alignment)
+            .with_skip_validation(skip_validation)
             .read_record_batch()?;
 
             Some(record_batch.column(0).clone())
@@ -772,6 +847,7 @@ pub struct FileDecoder {
     version: MetadataVersion,
     projection: Option<Vec<usize>>,
     require_alignment: bool,
+    skip_validation: UnsafeFlag,
 }
 
 impl FileDecoder {
@@ -783,6 +859,7 @@ impl FileDecoder {
             dictionaries: Default::default(),
             projection: None,
             require_alignment: false,
+            skip_validation: UnsafeFlag::new(),
         }
     }
 
@@ -809,6 +886,21 @@ impl FileDecoder {
         self
     }
 
+    /// Specifies whether validation should be skipped when reading data (default to `false`)
+    ///
+    /// # Safety
+    ///
+    /// This flag must only be set to `true` when you trust and are sure the data you are
+    /// reading is a valid Arrow IPC file, otherwise undefined behavior may
+    /// result.
+    ///
+    /// For example, some programs may wish to trust reading IPC files written
+    /// by the same process that created the files.
+    pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
+        self.skip_validation.set(skip_validation);
+        self
+    }
+
     fn read_message<'a>(&self, buf: &'a [u8]) -> Result<Message<'a>, ArrowError> {
         let message = parse_message(buf)?;
 
@@ -834,6 +926,7 @@ impl FileDecoder {
                     &mut self.dictionaries,
                     &message.version(),
                     self.require_alignment,
+                    self.skip_validation,
                 )
             }
             t => Err(ArrowError::ParseError(format!(
@@ -1250,6 +1343,9 @@ pub struct StreamReader<R> {
 
     /// Optional projection
     projection: Option<(Vec<usize>, Schema)>,
+
+    /// Should the reader skip validation
+    skip_validation: UnsafeFlag,
 }
 
 impl<R> fmt::Debug for StreamReader<R> {
@@ -1329,6 +1425,7 @@ impl<R: Read> StreamReader<R> {
             finished: false,
             dictionaries_by_id,
             projection,
+            skip_validation: UnsafeFlag::new(),
         })
     }
 
@@ -1437,6 +1534,7 @@ impl<R: Read> StreamReader<R> {
                     &mut self.dictionaries_by_id,
                     &message.version(),
                     false,
+                    self.skip_validation,
                 )?;
 
                 // read the next message until we encounter a RecordBatch
@@ -1462,6 +1560,21 @@ impl<R: Read> StreamReader<R> {
     pub fn get_mut(&mut self) -> &mut R {
         &mut self.reader
     }
+
+    /// Specifies whether validation should be skipped when reading data (default to `false`)
+    ///
+    /// # Safety
+    ///
+    /// This flag must only be set to `true` when you trust and are sure the data you are
+    /// reading is a valid Arrow IPC file, otherwise undefined behavior may
+    /// result.
+    ///
+    /// For example, some programs may wish to trust reading IPC files written
+    /// by the same process that created the files.
+    pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
+        self.skip_validation.set(skip_validation);
+        self
+    }
 }
 
 impl<R: Read> Iterator for StreamReader<R> {

arrow-ipc/src/reader/stream.rs

@@ -24,6 +24,7 @@ use arrow_buffer::{Buffer, MutableBuffer};
 use arrow_schema::{ArrowError, SchemaRef};
 
 use crate::convert::MessageBuffer;
+use crate::reader::private::UnsafeFlag;
 use crate::reader::{read_dictionary_impl, RecordBatchDecoder};
 use crate::{MessageHeader, CONTINUATION_MARKER};
 
@@ -42,6 +43,8 @@ pub struct StreamDecoder {
     buf: MutableBuffer,
     /// Whether or not array data in input buffers are required to be aligned
     require_alignment: bool,
+    /// Should we skip validation when reading arrays?
+    skip_validation: UnsafeFlag,
 }
 
 #[derive(Debug)]
@@ -102,6 +105,21 @@ impl StreamDecoder {
         self
     }
 
+    /// Specifies whether validation should be skipped when reading data (default to `false`)
+    ///
+    /// # Safety
+    ///
+    /// This flag must only be set to `true` when you trust and are sure the data you are
+    /// reading is a valid Arrow IPC stream, otherwise undefined behavior may
+    /// result.
+    ///
+    /// For example, some programs may wish to trust reading IPC streams written
+    /// by the same process that created the files.
+    pub unsafe fn with_skip_validation(mut self, skip_validation: bool) -> Self {
+        self.skip_validation.set(skip_validation);
+        self
+    }
+
     /// Try to read the next [`RecordBatch`] from the provided [`Buffer`]
     ///
     /// [`Buffer::advance`] will be called on `buffer` for any consumed bytes.
@@ -219,6 +237,7 @@ impl StreamDecoder {
                                 &version,
                             )?
                             .with_require_alignment(self.require_alignment)
+                            .with_skip_validation(self.skip_validation)
                             .read_record_batch()?;
                             self.state = DecoderState::default();
                             return Ok(Some(batch));
@@ -235,6 +254,7 @@ impl StreamDecoder {
                                 &mut self.dictionaries,
                                 &version,
                                 self.require_alignment,
+                                self.skip_validation,
                             )?;
                             self.state = DecoderState::default();
                         }