Alowed reusing compression buffer

Cargo.toml

@@ -14,7 +14,7 @@ bench = false
 [dependencies]
 parquet-format-async-temp = "0.2.0"
 bitpacking = { version = "0.8.2", features = ["bitpacker1x"] }
-streaming-iterator = "0.1.5"
+streaming-decompression = "0.1"
 
 async-stream = { version = "0.3.2", optional = true }
 futures = { version = "0.3", optional = true }

integration-tests/src/read/binary.rs

@@ -1,5 +1,5 @@
 use parquet::{
-    encoding::{bitpacking, hybrid_rle::HybridRleDecoder, plain_byte_array, uleb128, Encoding},
+    encoding::{hybrid_rle::HybridRleDecoder, plain_byte_array, Encoding},
     error::Result,
     metadata::ColumnDescriptor,
     page::{split_buffer, BinaryPageDict, DataPage},

integration-tests/src/read/mod.rs

@@ -21,55 +21,55 @@ pub fn page_to_array(page: &DataPage, descriptor: &ColumnDescriptor) -> Result<A
     match (descriptor.type_(), descriptor.max_rep_level()) {
         (ParquetType::PrimitiveType { physical_type, .. }, 0) => match page.dictionary_page() {
             Some(_) => match physical_type {
-                PhysicalType::Int32 => Ok(Array::Int32(primitive::page_dict_to_vec(
-                    &page, descriptor,
-                )?)),
-                PhysicalType::Int64 => Ok(Array::Int64(primitive::page_dict_to_vec(
-                    &page, descriptor,
-                )?)),
-                PhysicalType::Int96 => Ok(Array::Int96(primitive::page_dict_to_vec(
-                    &page, descriptor,
-                )?)),
+                PhysicalType::Int32 => {
+                    Ok(Array::Int32(primitive::page_dict_to_vec(page, descriptor)?))
+                }
+                PhysicalType::Int64 => {
+                    Ok(Array::Int64(primitive::page_dict_to_vec(page, descriptor)?))
+                }
+                PhysicalType::Int96 => {
+                    Ok(Array::Int96(primitive::page_dict_to_vec(page, descriptor)?))
+                }
                 PhysicalType::Float => Ok(Array::Float32(primitive::page_dict_to_vec(
-                    &page, descriptor,
+                    page, descriptor,
                 )?)),
                 PhysicalType::Double => Ok(Array::Float64(primitive::page_dict_to_vec(
-                    &page, descriptor,
+                    page, descriptor,
                 )?)),
                 PhysicalType::ByteArray => {
-                    Ok(Array::Binary(binary::page_dict_to_vec(&page, descriptor)?))
+                    Ok(Array::Binary(binary::page_dict_to_vec(page, descriptor)?))
                 }
                 _ => todo!(),
             },
             None => match physical_type {
                 PhysicalType::Boolean => {
-                    Ok(Array::Boolean(boolean::page_to_vec(&page, descriptor)?))
+                    Ok(Array::Boolean(boolean::page_to_vec(page, descriptor)?))
                 }
-                PhysicalType::Int32 => Ok(Array::Int32(primitive::page_to_vec(&page, descriptor)?)),
-                PhysicalType::Int64 => Ok(Array::Int64(primitive::page_to_vec(&page, descriptor)?)),
-                PhysicalType::Int96 => Ok(Array::Int96(primitive::page_to_vec(&page, descriptor)?)),
+                PhysicalType::Int32 => Ok(Array::Int32(primitive::page_to_vec(page, descriptor)?)),
+                PhysicalType::Int64 => Ok(Array::Int64(primitive::page_to_vec(page, descriptor)?)),
+                PhysicalType::Int96 => Ok(Array::Int96(primitive::page_to_vec(page, descriptor)?)),
                 PhysicalType::Float => {
-                    Ok(Array::Float32(primitive::page_to_vec(&page, descriptor)?))
+                    Ok(Array::Float32(primitive::page_to_vec(page, descriptor)?))
                 }
                 PhysicalType::Double => {
-                    Ok(Array::Float64(primitive::page_to_vec(&page, descriptor)?))
+                    Ok(Array::Float64(primitive::page_to_vec(page, descriptor)?))
                 }
                 PhysicalType::ByteArray => {
-                    Ok(Array::Binary(binary::page_to_vec(&page, descriptor)?))
+                    Ok(Array::Binary(binary::page_to_vec(page, descriptor)?))
                 }
                 _ => todo!(),
             },
         },
         (ParquetType::PrimitiveType { physical_type, .. }, _) => match page.dictionary_page() {
             None => match physical_type {
                 PhysicalType::Int64 => {
-                    Ok(primitive_nested::page_to_array::<i64>(&page, descriptor)?)
+                    Ok(primitive_nested::page_to_array::<i64>(page, descriptor)?)
                 }
                 _ => todo!(),
             },
             Some(_) => match physical_type {
                 PhysicalType::Int64 => Ok(primitive_nested::page_dict_to_array::<i64>(
-                    &page, descriptor,
+                    page, descriptor,
                 )?),
                 _ => todo!(),
             },
@@ -83,9 +83,10 @@ pub(crate) mod tests {
     use std::fs::File;
 
     use parquet::error::Result;
-    use parquet::read::{get_page_iterator, read_metadata, Decompressor, StreamingIterator};
+    use parquet::read::{get_page_iterator, read_metadata, Decompressor};
     use parquet::statistics::{BinaryStatistics, PrimitiveStatistics, Statistics};
     use parquet::types::int96_to_i64_ns;
+    use parquet::FallibleStreamingIterator;
 
     use super::*;
     use crate::tests::*;
@@ -100,14 +101,14 @@ pub(crate) mod tests {
         let column_meta = metadata.row_groups[row_group].column(column);
         let descriptor = column_meta.descriptor().clone();
 
-        let iterator = get_page_iterator(&column_meta, reader, None, vec![])?;
+        let iterator = get_page_iterator(column_meta, reader, None, vec![])?;
 
         let buffer = vec![];
         let mut iterator = Decompressor::new(iterator, buffer);
 
         let statistics = column_meta.statistics().transpose()?;
 
-        let page = iterator.next().unwrap().as_ref().unwrap();
+        let page = iterator.next()?.unwrap();
 
         let array = page_to_array(page, &descriptor)?;
 
@@ -201,7 +202,7 @@ pub(crate) mod tests {
         ];
 
         let expected = expected.into_iter().map(Some).collect::<Vec<_>>();
-        let (array, _) = get_column(&path, 10)?;
+        let (array, _) = get_column(path, 10)?;
         if let Array::Int96(array) = array {
             let a = array
                 .into_iter()

integration-tests/src/read/primitive.rs

@@ -3,7 +3,7 @@ use std::convert::TryInto;
 use super::utils::ValuesDef;
 
 use parquet::{
-    encoding::{bitpacking, hybrid_rle::HybridRleDecoder, uleb128, Encoding},
+    encoding::{hybrid_rle::HybridRleDecoder, Encoding},
     error::Result,
     metadata::ColumnDescriptor,
     page::{split_buffer, DataPage, PrimitivePageDict},

integration-tests/src/read/primitive_nested.rs

@@ -177,7 +177,7 @@ fn read_dict_array<T: NativeType>(
     let bit_width = values[0];
     let values = &values[1..];
 
-    let (_, consumed) = uleb128::decode(&values);
+    let (_, consumed) = uleb128::decode(values);
     let values = &values[consumed..];
 
     let indices = bitpacking::Decoder::new(values, bit_width, length as usize);

integration-tests/src/write/mod.rs

@@ -1,23 +1,21 @@
 pub(crate) mod primitive;
 
-use parquet::{
-    error::Result, metadata::ColumnDescriptor, page::CompressedPage, write::WriteOptions,
-};
+use parquet::{error::Result, metadata::ColumnDescriptor, page::EncodedPage, write::WriteOptions};
 
 use super::Array;
 
 pub fn array_to_page(
     array: &Array,
     options: &WriteOptions,
     descriptor: &ColumnDescriptor,
-) -> Result<CompressedPage> {
+) -> Result<EncodedPage> {
     // using plain encoding format
     match array {
-        Array::Int32(array) => primitive::array_to_page_v1(&array, options, descriptor),
-        Array::Int64(array) => primitive::array_to_page_v1(&array, options, descriptor),
-        Array::Int96(array) => primitive::array_to_page_v1(&array, options, descriptor),
-        Array::Float32(array) => primitive::array_to_page_v1(&array, options, descriptor),
-        Array::Float64(array) => primitive::array_to_page_v1(&array, options, descriptor),
+        Array::Int32(array) => primitive::array_to_page_v1(array, options, descriptor),
+        Array::Int64(array) => primitive::array_to_page_v1(array, options, descriptor),
+        Array::Int96(array) => primitive::array_to_page_v1(array, options, descriptor),
+        Array::Float32(array) => primitive::array_to_page_v1(array, options, descriptor),
+        Array::Float64(array) => primitive::array_to_page_v1(array, options, descriptor),
         _ => todo!(),
     }
 }
@@ -33,7 +31,7 @@ mod tests {
     use parquet::error::Result;
     use parquet::metadata::SchemaDescriptor;
     use parquet::statistics::Statistics;
-    use parquet::write::{write_file, DynIter, Version};
+    use parquet::write::{write_file, Compressor, DynIter, DynStreamingIterator, Version};
 
     use super::*;
 
@@ -67,9 +65,13 @@ mod tests {
 
         let a = schema.columns();
 
-        let row_groups = std::iter::once(Ok(DynIter::new(std::iter::once(Ok(DynIter::new(
-            std::iter::once(array_to_page(&array, &options, &a[0])),
-        ))))));
+        let row_groups = std::iter::once(Ok(DynIter::new(std::iter::once(Ok(
+            DynStreamingIterator::new(Compressor::new_from_vec(
+                DynIter::new(std::iter::once(array_to_page(&array, &options, &a[0]))),
+                options.compression,
+                vec![],
+            )),
+        )))));
 
         let mut writer = Cursor::new(vec![]);
         write_file(&mut writer, row_groups, schema, options, None, None)?;
@@ -140,7 +142,7 @@ mod tests2 {
         error::Result,
         metadata::SchemaDescriptor,
         read::read_metadata,
-        write::{write_file, DynIter, Version},
+        write::{write_file, Compressor, DynIter, DynStreamingIterator, Version},
     };
 
     #[test]
@@ -163,9 +165,17 @@ mod tests2 {
 
         let schema = SchemaDescriptor::try_from_message("message schema { OPTIONAL INT32 col; }")?;
 
-        let row_groups = std::iter::once(Ok(DynIter::new(std::iter::once(Ok(DynIter::new(
-            std::iter::once(array_to_page_v1(&array, &options, &schema.columns()[0])),
-        ))))));
+        let row_groups = std::iter::once(Ok(DynIter::new(std::iter::once(Ok(
+            DynStreamingIterator::new(Compressor::new_from_vec(
+                DynIter::new(std::iter::once(array_to_page_v1(
+                    &array,
+                    &options,
+                    &schema.columns()[0],
+                ))),
+                options.compression,
+                vec![],
+            )),
+        )))));
 
         let mut writer = Cursor::new(vec![]);
         write_file(&mut writer, row_groups, schema, options, None, None)?;

integration-tests/src/write/primitive.rs

@@ -1,11 +1,11 @@
 use parquet::{
     encoding::Encoding,
     metadata::ColumnDescriptor,
-    page::{CompressedDataPage, CompressedPage, DataPageHeader, DataPageHeaderV1},
+    page::{DataPage, DataPageHeader, DataPageHeaderV1, EncodedPage},
     statistics::{serialize_statistics, PrimitiveStatistics, Statistics},
     types::NativeType,
     write::WriteOptions,
-    {compression::create_codec, encoding::hybrid_rle::encode_bool, error::Result},
+    {encoding::hybrid_rle::encode_bool, error::Result},
 };
 
 fn unzip_option<T: NativeType>(array: &[Option<T>]) -> Result<(Vec<u8>, Vec<u8>)> {
@@ -44,22 +44,10 @@ pub fn array_to_page_v1<T: NativeType>(
     array: &[Option<T>],
     options: &WriteOptions,
     descriptor: &ColumnDescriptor,
-) -> Result<CompressedPage> {
+) -> Result<EncodedPage> {
     let (values, mut buffer) = unzip_option(array)?;
 
     buffer.extend_from_slice(&values);
-    let uncompressed_page_size = buffer.len();
-
-    let codec = create_codec(&options.compression)?;
-    let buffer = if let Some(mut codec) = codec {
-        // todo: remove this allocation by extending `buffer` directly.
-        // needs refactoring `compress`'s API.
-        let mut tmp = vec![];
-        codec.compress(&values, &mut tmp)?;
-        tmp
-    } else {
-        buffer
-    };
 
     let statistics = if options.write_statistics {
         let statistics = &PrimitiveStatistics {
@@ -82,11 +70,9 @@ pub fn array_to_page_v1<T: NativeType>(
         statistics,
     };
 
-    Ok(CompressedPage::Data(CompressedDataPage::new(
+    Ok(EncodedPage::Data(DataPage::new(
         DataPageHeader::V1(header),
         buffer,
-        options.compression,
-        uncompressed_page_size,
         None,
         descriptor.clone(),
     )))

src/compression.rs

@@ -274,33 +274,25 @@ mod tests {
 
     fn test_roundtrip(c: Compression, data: &[u8]) {
         let mut c1 = create_codec(&c).unwrap().unwrap();
-        let mut c2 = create_codec(&c).unwrap().unwrap();
 
-        // Compress with c1
-        let mut compressed = Vec::new();
+        let offset = 2;
+
+        // Compress to a buffer that already has data is possible
+        let mut compressed = vec![2; offset];
         c1.compress(data, &mut compressed)
             .expect("Error when compressing");
 
-        // Decompress with c2
-        let mut decompressed = vec![0; data.len()];
-        c2.decompress(compressed.as_slice(), &mut decompressed)
-            .expect("Error when decompressing");
-        assert_eq!(data, decompressed.as_slice());
-
-        compressed.clear();
-
-        // Compress with c2
-        c2.compress(data, &mut compressed)
-            .expect("Error when compressing");
+        // data is compressed...
+        assert!(compressed.len() - 2 < data.len());
 
-        // Decompress with c1
-        c1.decompress(compressed.as_slice(), &mut decompressed)
+        let mut decompressed = vec![0; data.len()];
+        c1.decompress(&compressed[offset..], &mut decompressed)
             .expect("Error when decompressing");
         assert_eq!(data, decompressed.as_slice());
     }
 
     fn test_codec(c: Compression) {
-        let sizes = vec![100, 10000, 100000];
+        let sizes = vec![10000, 100000];
         for size in sizes {
             let data = (0..size).map(|x| (x % 255) as u8).collect::<Vec<_>>();
             test_roundtrip(c, &data);

src/error.rs

@@ -28,13 +28,6 @@ impl std::fmt::Display for ParquetError {
     }
 }
 
-impl ParquetError {
-    /// Wraps an external error in an `ParquetError`.
-    pub fn from_external_error(error: impl std::error::Error + Send + Sync + 'static) -> Self {
-        Self::External("".to_string(), Arc::new(error))
-    }
-}
-
 #[cfg(feature = "snappy")]
 impl From<snap::Error> for ParquetError {
     fn from(e: snap::Error) -> ParquetError {

src/lib.rs

@@ -14,6 +14,9 @@ pub mod statistics;
 pub mod types;
 pub mod write;
 
+pub use streaming_decompression::fallible_streaming_iterator;
+pub use streaming_decompression::FallibleStreamingIterator;
+
 const FOOTER_SIZE: u64 = 8;
 const PARQUET_MAGIC: [u8; 4] = [b'P', b'A', b'R', b'1'];