make sure that scratch is used in decompressing data

jorgecarleitao · Jun 27, 2022 · 5bed494 · 5bed494
1 parent 53d1b54
commit 5bed494
Show file tree

Hide file tree

Showing 19 changed files with 127 additions and 22 deletions.
diff --git a/examples/ipc_file_read.rs b/examples/ipc_file_read.rs
@@ -33,7 +33,7 @@ fn read_batch(path: &str) -> Result<(Schema, Chunk<Box<dyn Array>>)> {
     let schema = metadata.schema.clone();
 
     // advanced way: read the dictionary
-    let dictionaries = read::read_file_dictionaries(&mut file, &metadata, Default::default())?;
+    let dictionaries = read::read_file_dictionaries(&mut file, &metadata, &mut Default::default())?;
 
     let chunk_index = 0;
 
@@ -44,6 +44,7 @@ fn read_batch(path: &str) -> Result<(Schema, Chunk<Box<dyn Array>>)> {
         None,
         chunk_index,
         &mut Default::default(),
+        &mut Default::default(),
     )?;
 
     Ok((schema, chunk))

diff --git a/src/io/flight/mod.rs b/src/io/flight/mod.rs
@@ -142,6 +142,7 @@ pub fn deserialize_batch(
             &mut reader,
             0,
             length as u64,
+            &mut Default::default(),
         ),
         _ => Err(Error::nyi(
             "flight currently only supports reading RecordBatch messages",

diff --git a/src/io/ipc/read/array/binary.rs b/src/io/ipc/read/array/binary.rs
@@ -5,10 +5,12 @@ use crate::array::{BinaryArray, Offset};
 use crate::buffer::Buffer;
 use crate::datatypes::DataType;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::read_basic::*;
 use super::super::{Compression, IpcBuffer, Node, OutOfSpecKind};
 
+#[allow(clippy::too_many_arguments)]
 pub fn read_binary<O: Offset, R: Read + Seek>(
     field_nodes: &mut VecDeque<Node>,
     data_type: DataType,
@@ -17,6 +19,7 @@ pub fn read_binary<O: Offset, R: Read + Seek>(
     block_offset: u64,
     is_little_endian: bool,
     compression: Option<Compression>,
+    scratch: &mut ReadBuffer,
 ) -> Result<BinaryArray<O>> {
     let field_node = field_nodes.pop_front().ok_or_else(|| {
         Error::oos(format!(
@@ -46,6 +49,7 @@ pub fn read_binary<O: Offset, R: Read + Seek>(
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )
     // Older versions of the IPC format sometimes do not report an offset
     .or_else(|_| Result::Ok(Buffer::<O>::from(vec![O::default()])))?;
@@ -58,6 +62,7 @@ pub fn read_binary<O: Offset, R: Read + Seek>(
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )?;
 
     BinaryArray::<O>::try_new(data_type, offsets, values, validity)

diff --git a/src/io/ipc/read/array/dictionary.rs b/src/io/ipc/read/array/dictionary.rs
@@ -4,6 +4,7 @@ use std::io::{Read, Seek};
 
 use crate::array::{DictionaryArray, DictionaryKey};
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::Dictionaries;
 use super::super::{Compression, IpcBuffer, Node};
@@ -19,6 +20,7 @@ pub fn read_dictionary<T: DictionaryKey, R: Read + Seek>(
     block_offset: u64,
     compression: Option<Compression>,
     is_little_endian: bool,
+    scratch: &mut ReadBuffer,
 ) -> Result<DictionaryArray<T>>
 where
     Vec<u8>: TryInto<T::Bytes>,
@@ -47,6 +49,7 @@ where
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )?;
 
     Ok(DictionaryArray::<T>::from_data(keys, values))

diff --git a/src/io/ipc/read/array/fixed_size_binary.rs b/src/io/ipc/read/array/fixed_size_binary.rs
@@ -4,10 +4,12 @@ use std::io::{Read, Seek};
 use crate::array::FixedSizeBinaryArray;
 use crate::datatypes::DataType;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::read_basic::*;
 use super::super::{Compression, IpcBuffer, Node, OutOfSpecKind};
 
+#[allow(clippy::too_many_arguments)]
 pub fn read_fixed_size_binary<R: Read + Seek>(
     field_nodes: &mut VecDeque<Node>,
     data_type: DataType,
@@ -16,6 +18,7 @@ pub fn read_fixed_size_binary<R: Read + Seek>(
     block_offset: u64,
     is_little_endian: bool,
     compression: Option<Compression>,
+    scratch: &mut ReadBuffer,
 ) -> Result<FixedSizeBinaryArray> {
     let field_node = field_nodes.pop_front().ok_or_else(|| {
         Error::oos(format!(
@@ -46,6 +49,7 @@ pub fn read_fixed_size_binary<R: Read + Seek>(
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )?;
 
     FixedSizeBinaryArray::try_new(data_type, values, validity)

diff --git a/src/io/ipc/read/array/fixed_size_list.rs b/src/io/ipc/read/array/fixed_size_list.rs
@@ -4,6 +4,7 @@ use std::io::{Read, Seek};
 use crate::array::FixedSizeListArray;
 use crate::datatypes::DataType;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::super::IpcField;
 use super::super::deserialize::{read, skip};
@@ -23,6 +24,7 @@ pub fn read_fixed_size_list<R: Read + Seek>(
     is_little_endian: bool,
     compression: Option<Compression>,
     version: Version,
+    scratch: &mut ReadBuffer,
 ) -> Result<FixedSizeListArray> {
     let field_node = field_nodes.pop_front().ok_or_else(|| {
         Error::oos(format!(
@@ -53,6 +55,7 @@ pub fn read_fixed_size_list<R: Read + Seek>(
         is_little_endian,
         compression,
         version,
+        scratch,
     )?;
     FixedSizeListArray::try_new(data_type, values, validity)
 }

diff --git a/src/io/ipc/read/array/list.rs b/src/io/ipc/read/array/list.rs
@@ -6,6 +6,7 @@ use crate::array::{ListArray, Offset};
 use crate::buffer::Buffer;
 use crate::datatypes::DataType;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::super::IpcField;
 use super::super::deserialize::{read, skip};
@@ -25,6 +26,7 @@ pub fn read_list<O: Offset, R: Read + Seek>(
     is_little_endian: bool,
     compression: Option<Compression>,
     version: Version,
+    scratch: &mut ReadBuffer,
 ) -> Result<ListArray<O>>
 where
     Vec<u8>: TryInto<O::Bytes>,
@@ -57,6 +59,7 @@ where
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )
     // Older versions of the IPC format sometimes do not report an offset
     .or_else(|_| Result::Ok(Buffer::<O>::from(vec![O::default()])))?;
@@ -74,6 +77,7 @@ where
         is_little_endian,
         compression,
         version,
+        scratch,
     )?;
     ListArray::try_new(data_type, offsets, values, validity)
 }

diff --git a/src/io/ipc/read/array/map.rs b/src/io/ipc/read/array/map.rs
@@ -5,6 +5,7 @@ use crate::array::MapArray;
 use crate::buffer::Buffer;
 use crate::datatypes::DataType;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::super::IpcField;
 use super::super::deserialize::{read, skip};
@@ -24,6 +25,7 @@ pub fn read_map<R: Read + Seek>(
     is_little_endian: bool,
     compression: Option<Compression>,
     version: Version,
+    scratch: &mut ReadBuffer,
 ) -> Result<MapArray> {
     let field_node = field_nodes.pop_front().ok_or_else(|| {
         Error::oos(format!(
@@ -53,6 +55,7 @@ pub fn read_map<R: Read + Seek>(
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )
     // Older versions of the IPC format sometimes do not report an offset
     .or_else(|_| Result::Ok(Buffer::<i32>::from(vec![0i32])))?;
@@ -70,6 +73,7 @@ pub fn read_map<R: Read + Seek>(
         is_little_endian,
         compression,
         version,
+        scratch,
     )?;
     MapArray::try_new(data_type, offsets, field, validity)
 }

diff --git a/src/io/ipc/read/array/primitive.rs b/src/io/ipc/read/array/primitive.rs
@@ -3,11 +3,13 @@ use std::{collections::VecDeque, convert::TryInto};
 
 use crate::datatypes::DataType;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 use crate::{array::PrimitiveArray, types::NativeType};
 
 use super::super::read_basic::*;
 use super::super::{Compression, IpcBuffer, Node, OutOfSpecKind};
 
+#[allow(clippy::too_many_arguments)]
 pub fn read_primitive<T: NativeType, R: Read + Seek>(
     field_nodes: &mut VecDeque<Node>,
     data_type: DataType,
@@ -16,6 +18,7 @@ pub fn read_primitive<T: NativeType, R: Read + Seek>(
     block_offset: u64,
     is_little_endian: bool,
     compression: Option<Compression>,
+    scratch: &mut ReadBuffer,
 ) -> Result<PrimitiveArray<T>>
 where
     Vec<u8>: TryInto<T::Bytes>,
@@ -48,6 +51,7 @@ where
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )?;
     PrimitiveArray::<T>::try_new(data_type, values, validity)
 }

diff --git a/src/io/ipc/read/array/struct_.rs b/src/io/ipc/read/array/struct_.rs
@@ -4,6 +4,7 @@ use std::io::{Read, Seek};
 use crate::array::StructArray;
 use crate::datatypes::DataType;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::super::IpcField;
 use super::super::deserialize::{read, skip};
@@ -23,6 +24,7 @@ pub fn read_struct<R: Read + Seek>(
     is_little_endian: bool,
     compression: Option<Compression>,
     version: Version,
+    scratch: &mut ReadBuffer,
 ) -> Result<StructArray> {
     let field_node = field_nodes.pop_front().ok_or_else(|| {
         Error::oos(format!(
@@ -57,6 +59,7 @@ pub fn read_struct<R: Read + Seek>(
                 is_little_endian,
                 compression,
                 version,
+                scratch,
             )
         })
         .collect::<Result<Vec<_>>>()?;

diff --git a/src/io/ipc/read/array/union.rs b/src/io/ipc/read/array/union.rs
@@ -5,6 +5,7 @@ use crate::array::UnionArray;
 use crate::datatypes::DataType;
 use crate::datatypes::UnionMode::Dense;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::super::IpcField;
 use super::super::deserialize::{read, skip};
@@ -24,6 +25,7 @@ pub fn read_union<R: Read + Seek>(
     is_little_endian: bool,
     compression: Option<Compression>,
     version: Version,
+    scratch: &mut ReadBuffer,
 ) -> Result<UnionArray> {
     let field_node = field_nodes.pop_front().ok_or_else(|| {
         Error::oos(format!(
@@ -50,6 +52,7 @@ pub fn read_union<R: Read + Seek>(
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )?;
 
     let offsets = if let DataType::Union(_, _, mode) = data_type {
@@ -61,6 +64,7 @@ pub fn read_union<R: Read + Seek>(
                 block_offset,
                 is_little_endian,
                 compression,
+                scratch,
             )?)
         } else {
             None
@@ -86,6 +90,7 @@ pub fn read_union<R: Read + Seek>(
                 is_little_endian,
                 compression,
                 version,
+                scratch,
             )
         })
         .collect::<Result<Vec<_>>>()?;

diff --git a/src/io/ipc/read/array/utf8.rs b/src/io/ipc/read/array/utf8.rs
@@ -5,10 +5,12 @@ use crate::array::{Offset, Utf8Array};
 use crate::buffer::Buffer;
 use crate::datatypes::DataType;
 use crate::error::{Error, Result};
+use crate::io::ipc::read::common::ReadBuffer;
 
 use super::super::read_basic::*;
 use super::super::{Compression, IpcBuffer, Node, OutOfSpecKind};
 
+#[allow(clippy::too_many_arguments)]
 pub fn read_utf8<O: Offset, R: Read + Seek>(
     field_nodes: &mut VecDeque<Node>,
     data_type: DataType,
@@ -17,6 +19,7 @@ pub fn read_utf8<O: Offset, R: Read + Seek>(
     block_offset: u64,
     is_little_endian: bool,
     compression: Option<Compression>,
+    scratch: &mut ReadBuffer,
 ) -> Result<Utf8Array<O>> {
     let field_node = field_nodes.pop_front().ok_or_else(|| {
         Error::oos(format!(
@@ -46,6 +49,7 @@ pub fn read_utf8<O: Offset, R: Read + Seek>(
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )
     // Older versions of the IPC format sometimes do not report an offset
     .or_else(|_| Result::Ok(Buffer::<O>::from(vec![O::default()])))?;
@@ -58,6 +62,7 @@ pub fn read_utf8<O: Offset, R: Read + Seek>(
         block_offset,
         is_little_endian,
         compression,
+        scratch,
     )?;
 
     Utf8Array::<O>::try_new(data_type, offsets, values, validity)

diff --git a/src/io/ipc/read/common.rs b/src/io/ipc/read/common.rs
@@ -85,6 +85,7 @@ pub fn read_record_batch<R: Read + Seek>(
     reader: &mut R,
     block_offset: u64,
     file_size: u64,
+    scratch: &mut ReadBuffer,
 ) -> Result<Chunk<Box<dyn Array>>> {
     assert_eq!(fields.len(), ipc_schema.fields.len());
     let buffers = batch
@@ -136,6 +137,7 @@ pub fn read_record_batch<R: Read + Seek>(
                         Error::from(OutOfSpecKind::InvalidFlatbufferCompression(err))
                     })?,
                     version,
+                    scratch,
                 )?)),
                 ProjectionResult::NotSelected((field, _)) => {
                     skip(&mut field_nodes, &field.data_type, &mut buffers)?;
@@ -162,6 +164,7 @@ pub fn read_record_batch<R: Read + Seek>(
                         Error::from(OutOfSpecKind::InvalidFlatbufferCompression(err))
                     })?,
                     version,
+                    scratch,
                 )
             })
             .collect::<Result<Vec<_>>>()?
@@ -221,6 +224,7 @@ fn first_dict_field<'a>(
 
 /// Read the dictionary from the buffer and provided metadata,
 /// updating the `dictionaries` with the resulting dictionary
+#[allow(clippy::too_many_arguments)]
 pub fn read_dictionary<R: Read + Seek>(
     batch: arrow_format::ipc::DictionaryBatchRef,
     fields: &[Field],
@@ -229,6 +233,7 @@ pub fn read_dictionary<R: Read + Seek>(
     reader: &mut R,
     block_offset: u64,
     file_size: u64,
+    scratch: &mut ReadBuffer,
 ) -> Result<()> {
     if batch
         .is_delta()
@@ -270,6 +275,7 @@ pub fn read_dictionary<R: Read + Seek>(
                 reader,
                 block_offset,
                 file_size,
+                scratch,
             )?;
             let mut arrays = columns.into_arrays();
             arrays.pop().unwrap()