Added support to read parquet asynchronously.

jorgecarleitao · Aug 9, 2021 · 1ada073 · 1ada073
1 parent 988612a
commit 1ada073
Show file tree

Hide file tree

Showing 28 changed files with 125 additions and 105 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -59,7 +59,7 @@ futures = { version = "0.3", optional = true }
 ahash = { version = "0.7", optional = true }
 
 #parquet2 = { version = "0.2", optional = true, default_features = false, features = ["stream"] }
-parquet2 = { git = "https://github.com/jorgecarleitao/parquet2", ref = "65990906a90edf93530a2c7a54b31b023974e5f3", optional = true, default_features = false, features = ["stream"] }
+parquet2 = { git = "https://github.com/jorgecarleitao/parquet2", rev = "41ca6b3fdbe62f9fa5f1758748dce9a8db4f72d6", optional = true, default_features = false, features = ["stream"] }
 
 [dev-dependencies]
 rand = "0.8"

diff --git a/arrow-parquet-integration-testing/src/main.rs b/arrow-parquet-integration-testing/src/main.rs
@@ -7,7 +7,7 @@ use arrow2::error::Result;
 use arrow2::io::parquet::write::{Encoding, RowGroupIterator};
 use arrow2::io::{
     json_integration::ArrowJson,
-    parquet::write::{write_file, CompressionCodec, Version, WriteOptions},
+    parquet::write::{write_file, Compression, Version, WriteOptions},
 };
 use arrow2::{datatypes::Schema, io::json_integration::to_record_batch, record_batch::RecordBatch};
 
@@ -154,7 +154,7 @@ fn main() -> Result<()> {
 
     let options = WriteOptions {
         write_statistics: true,
-        compression: CompressionCodec::Uncompressed,
+        compression: Compression::Uncompressed,
         version,
     };
 

diff --git a/examples/parquet_write.rs b/examples/parquet_write.rs
@@ -7,7 +7,7 @@ use arrow2::{
     datatypes::{Field, Schema},
     error::Result,
     io::parquet::write::{
-        array_to_page, write_file, CompressionCodec, DynIter, Encoding, Version, WriteOptions,
+        array_to_page, write_file, Compression, DynIter, Encoding, Version, WriteOptions,
     },
 };
 
@@ -16,7 +16,7 @@ fn write_single_array(path: &str, array: &dyn Array, field: Field) -> Result<()>
 
     let options = WriteOptions {
         write_statistics: true,
-        compression: CompressionCodec::Uncompressed,
+        compression: Compression::Uncompressed,
         version: Version::V2,
     };
     let encoding = Encoding::Plain;

diff --git a/examples/parquet_write_record.rs b/examples/parquet_write_record.rs
@@ -5,17 +5,18 @@ use arrow2::{
     array::{Array, Int32Array},
     datatypes::{Field, Schema},
     error::Result,
-    io::parquet::write::{write_file, CompressionCodec, RowGroupIterator, Version, WriteOptions},
+    io::parquet::write::{
+        write_file, Compression, Encoding, RowGroupIterator, Version, WriteOptions,
+    },
     record_batch::RecordBatch,
 };
-use parquet2::schema::Encoding;
 
 fn write_batch(path: &str, batch: RecordBatch) -> Result<()> {
     let schema = batch.schema().clone();
 
     let options = WriteOptions {
         write_statistics: true,
-        compression: CompressionCodec::Uncompressed,
+        compression: Compression::Uncompressed,
         version: Version::V2,
     };
 

diff --git a/src/io/parquet/mod.rs b/src/io/parquet/mod.rs
@@ -405,7 +405,7 @@ mod tests {
 mod tests_integration {
     use std::sync::Arc;
 
-    use super::write::CompressionCodec;
+    use super::write::Compression;
     use crate::array::{Array, PrimitiveArray, Utf8Array};
     use crate::datatypes::DataType;
     use crate::datatypes::TimeUnit;
@@ -421,7 +421,7 @@ mod tests_integration {
     fn integration_write(schema: &Schema, batches: &[RecordBatch]) -> Result<Vec<u8>> {
         let options = WriteOptions {
             write_statistics: true,
-            compression: CompressionCodec::Uncompressed,
+            compression: Compression::Uncompressed,
             version: Version::V1,
         };
 

diff --git a/src/io/parquet/read/binary/basic.rs b/src/io/parquet/read/binary/basic.rs
@@ -2,7 +2,7 @@ use futures::{pin_mut, Stream, StreamExt};
 use parquet2::{
     encoding::{bitpacking, delta_length_byte_array, hybrid_rle, uleb128, Encoding},
     metadata::{ColumnChunkMetaData, ColumnDescriptor},
-    page::{BinaryPageDict, DataPage, DataPageHeader},
+    page::{BinaryPageDict, DataPage, DataPageHeader, DataPageHeaderExt},
     read::{levels, StreamingIterator},
 };
 
@@ -215,7 +215,7 @@ fn extend_from_page<O: Offset>(
     let is_optional = descriptor.max_def_level() == 1;
     match page.header() {
         DataPageHeader::V1(header) => {
-            assert_eq!(header.definition_level_encoding, Encoding::Rle);
+            assert_eq!(header.definition_level_encoding(), Encoding::Rle);
 
             let (_, validity_buffer, values_buffer) =
                 levels::split_buffer_v1(page.buffer(), false, is_optional);

diff --git a/src/io/parquet/read/binary/nested.rs b/src/io/parquet/read/binary/nested.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use parquet2::{
     encoding::Encoding,
     metadata::{ColumnChunkMetaData, ColumnDescriptor},
-    page::{DataPage, DataPageHeader},
+    page::{DataPage, DataPageHeader, DataPageHeaderExt},
     read::{
         levels::{get_bit_width, split_buffer_v1, split_buffer_v2, RLEDecoder},
         StreamingIterator,
@@ -121,8 +121,8 @@ fn extend_from_page<O: Offset>(
 
     match page.header() {
         DataPageHeader::V1(header) => {
-            assert_eq!(header.definition_level_encoding, Encoding::Rle);
-            assert_eq!(header.repetition_level_encoding, Encoding::Rle);
+            assert_eq!(header.definition_level_encoding(), Encoding::Rle);
+            assert_eq!(header.repetition_level_encoding(), Encoding::Rle);
 
             match (&page.encoding(), page.dictionary_page()) {
                 (Encoding::Plain, None) => {
@@ -137,11 +137,11 @@ fn extend_from_page<O: Offset>(
                         values_buffer,
                         additional,
                         (
-                            &header.repetition_level_encoding,
+                            &header.repetition_level_encoding(),
                             descriptor.max_rep_level(),
                         ),
                         (
-                            &header.definition_level_encoding,
+                            &header.definition_level_encoding(),
                             descriptor.max_def_level(),
                         ),
                         is_nullable,

diff --git a/src/io/parquet/read/boolean/basic.rs b/src/io/parquet/read/boolean/basic.rs
@@ -10,7 +10,7 @@ use futures::{pin_mut, Stream, StreamExt};
 use parquet2::{
     encoding::{hybrid_rle, Encoding},
     metadata::{ColumnChunkMetaData, ColumnDescriptor},
-    page::{DataPage, DataPageHeader},
+    page::{DataPage, DataPageHeader, DataPageHeaderExt},
     read::{levels, StreamingIterator},
 };
 
@@ -123,7 +123,7 @@ fn extend_from_page(
     let is_optional = descriptor.max_def_level() == 1;
     match page.header() {
         DataPageHeader::V1(header) => {
-            assert_eq!(header.definition_level_encoding, Encoding::Rle);
+            assert_eq!(header.definition_level_encoding(), Encoding::Rle);
 
             match (&page.encoding(), page.dictionary_page(), is_optional) {
                 (Encoding::Plain, None, true) => {

diff --git a/src/io/parquet/read/boolean/nested.rs b/src/io/parquet/read/boolean/nested.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use parquet2::{
     encoding::Encoding,
     metadata::{ColumnChunkMetaData, ColumnDescriptor},
-    page::{DataPage, DataPageHeader},
+    page::{DataPage, DataPageHeader, DataPageHeaderExt},
     read::{
         levels::{get_bit_width, split_buffer_v1, split_buffer_v2, RLEDecoder},
         StreamingIterator,
@@ -107,8 +107,8 @@ fn extend_from_page(
 
     match page.header() {
         DataPageHeader::V1(header) => {
-            assert_eq!(header.definition_level_encoding, Encoding::Rle);
-            assert_eq!(header.repetition_level_encoding, Encoding::Rle);
+            assert_eq!(header.definition_level_encoding(), Encoding::Rle);
+            assert_eq!(header.repetition_level_encoding(), Encoding::Rle);
 
             match (&page.encoding(), page.dictionary_page()) {
                 (Encoding::Plain, None) => {
@@ -123,11 +123,11 @@ fn extend_from_page(
                         values_buffer,
                         additional,
                         (
-                            &header.repetition_level_encoding,
+                            &header.repetition_level_encoding(),
                             descriptor.max_rep_level(),
                         ),
                         (
-                            &header.definition_level_encoding,
+                            &header.definition_level_encoding(),
                             descriptor.max_def_level(),
                         ),
                         is_nullable,

diff --git a/src/io/parquet/read/fixed_size_binary.rs b/src/io/parquet/read/fixed_size_binary.rs
@@ -1,6 +1,6 @@
 use parquet2::{
     encoding::{bitpacking, hybrid_rle, uleb128, Encoding},
-    page::{DataPage, DataPageHeader, FixedLenByteArrayPageDict},
+    page::{DataPage, DataPageHeader, DataPageHeaderExt, FixedLenByteArrayPageDict},
     read::{levels, StreamingIterator},
 };
 
@@ -171,7 +171,7 @@ pub(crate) fn extend_from_page(
     let is_optional = descriptor.max_def_level() == 1;
     match page.header() {
         DataPageHeader::V1(header) => {
-            assert_eq!(header.definition_level_encoding, Encoding::Rle);
+            assert_eq!(header.definition_level_encoding(), Encoding::Rle);
 
             let (_, validity_buffer, values_buffer) =
                 levels::split_buffer_v1(page.buffer(), false, is_optional);

diff --git a/src/io/parquet/read/mod.rs b/src/io/parquet/read/mod.rs
@@ -1,16 +1,17 @@
 use std::io::{Read, Seek};
 
-use futures::Stream;
+use futures::{AsyncRead, AsyncSeek, Stream};
 pub use parquet2::{
     error::ParquetError,
     metadata::{ColumnChunkMetaData, ColumnDescriptor, RowGroupMetaData},
     page::{CompressedDataPage, DataPage, DataPageHeader},
     read::{
-        decompress, get_page_iterator as _get_page_iterator, read_metadata as _read_metadata,
+        decompress, get_page_iterator as _get_page_iterator, get_page_stream as _get_page_stream,
+        read_metadata as _read_metadata, read_metadata_async as _read_metadata_async,
         streaming_iterator, Decompressor, PageIterator, StreamingIterator,
     },
-    schema::{
-        types::{LogicalType, ParquetType, PhysicalType, PrimitiveConvertedType},
+    schema::types::{
+        LogicalType, ParquetType, PhysicalType, PrimitiveConvertedType,
         TimeUnit as ParquetTimeUnit, TimestampType,
     },
     types::int96_to_i64_ns,
@@ -48,11 +49,29 @@ pub fn get_page_iterator<'b, RR: Read + Seek>(
     )?)
 }
 
-/// Reads parquets' metadata.
+/// Creates a new iterator of compressed pages.
+pub async fn get_page_stream<'a, RR: AsyncRead + Unpin + Send + AsyncSeek>(
+    metadata: &'a FileMetaData,
+    row_group: usize,
+    column: usize,
+    reader: &'a mut RR,
+    buffer: Vec<u8>,
+) -> Result<impl Stream<Item = std::result::Result<CompressedDataPage, ParquetError>> + 'a> {
+    Ok(_get_page_stream(metadata, row_group, column, reader, buffer).await?)
+}
+
+/// Reads parquets' metadata syncronously.
 pub fn read_metadata<R: Read + Seek>(reader: &mut R) -> Result<FileMetaData> {
     Ok(_read_metadata(reader)?)
 }
 
+/// Reads parquets' metadata asynchronously.
+pub async fn read_metadata_async<R: AsyncRead + AsyncSeek + Send + Unpin>(
+    reader: &mut R,
+) -> Result<FileMetaData> {
+    Ok(_read_metadata_async(reader).await?)
+}
+
 pub fn page_iter_to_array<
     I: StreamingIterator<Item = std::result::Result<DataPage, ParquetError>>,
 >(

diff --git a/src/io/parquet/read/primitive/basic.rs b/src/io/parquet/read/primitive/basic.rs
@@ -1,6 +1,6 @@
 use parquet2::{
     encoding::{bitpacking, hybrid_rle, uleb128, Encoding},
-    page::{DataPage, DataPageHeader, PrimitivePageDict},
+    page::{DataPage, DataPageHeader, DataPageHeaderExt, PrimitivePageDict},
     read::levels,
     types::NativeType,
 };
@@ -160,7 +160,7 @@ where
     let is_optional = descriptor.max_def_level() == 1;
     match page.header() {
         DataPageHeader::V1(header) => {
-            assert_eq!(header.definition_level_encoding, Encoding::Rle);
+            assert_eq!(header.definition_level_encoding(), Encoding::Rle);
 
             let (_, validity_buffer, values_buffer) =
                 levels::split_buffer_v1(page.buffer(), false, is_optional);

diff --git a/src/io/parquet/read/primitive/dictionary.rs b/src/io/parquet/read/primitive/dictionary.rs
@@ -1,10 +1,9 @@
 use std::sync::Arc;
 
 use parquet2::{
-    encoding::{bitpacking, hybrid_rle, uleb128},
-    page::{DataPage, DataPageHeader, PrimitivePageDict},
+    encoding::{bitpacking, hybrid_rle, uleb128, Encoding},
+    page::{DataPage, DataPageHeader, DataPageHeaderExt, PrimitivePageDict},
     read::{levels, StreamingIterator},
-    schema::Encoding,
     types::NativeType,
 };
 
@@ -102,7 +101,7 @@ where
     let is_optional = descriptor.max_def_level() == 1;
     match page.header() {
         DataPageHeader::V1(header) => {
-            assert_eq!(header.definition_level_encoding, Encoding::Rle);
+            assert_eq!(header.definition_level_encoding(), Encoding::Rle);
 
             let (_, validity_buffer, values_buffer) =
                 levels::split_buffer_v1(page.buffer(), false, is_optional);

diff --git a/src/io/parquet/read/primitive/nested.rs b/src/io/parquet/read/primitive/nested.rs
@@ -1,6 +1,6 @@
 use parquet2::{
     encoding::Encoding,
-    page::{DataPage, DataPageHeader},
+    page::{DataPage, DataPageHeader, DataPageHeaderExt},
     read::levels::{get_bit_width, split_buffer_v1, split_buffer_v2, RLEDecoder},
     types::NativeType,
 };
@@ -127,8 +127,8 @@ where
 
     match page.header() {
         DataPageHeader::V1(header) => {
-            assert_eq!(header.definition_level_encoding, Encoding::Rle);
-            assert_eq!(header.repetition_level_encoding, Encoding::Rle);
+            assert_eq!(header.definition_level_encoding(), Encoding::Rle);
+            assert_eq!(header.repetition_level_encoding(), Encoding::Rle);
 
             match (&page.encoding(), page.dictionary_page()) {
                 (Encoding::Plain, None) => {
@@ -143,11 +143,11 @@ where
                         values_buffer,
                         additional,
                         (
-                            &header.repetition_level_encoding,
+                            &header.repetition_level_encoding(),
                             descriptor.max_rep_level(),
                         ),
                         (
-                            &header.definition_level_encoding,
+                            &header.definition_level_encoding(),
                             descriptor.max_def_level(),
                         ),
                         is_nullable,

diff --git a/src/io/parquet/read/schema/convert.rs b/src/io/parquet/read/schema/convert.rs
@@ -6,9 +6,9 @@ use parquet2::{
     schema::{
         types::{
             BasicTypeInfo, GroupConvertedType, LogicalType, ParquetType, PhysicalType,
-            PrimitiveConvertedType, TimeUnit as ParquetTimeUnit,
+            PrimitiveConvertedType, TimeUnit as ParquetTimeUnit, TimestampType,
         },
-        Repetition, TimestampType,
+        Repetition,
     },
 };
 

diff --git a/src/io/parquet/read/utils.rs b/src/io/parquet/read/utils.rs
@@ -1,4 +1,4 @@
-use parquet2::{encoding::get_length, schema::Encoding};
+use parquet2::encoding::{get_length, Encoding};
 
 use crate::error::ArrowError;
 

diff --git a/src/io/parquet/write/binary/nested.rs b/src/io/parquet/write/binary/nested.rs
@@ -1,4 +1,4 @@
-use parquet2::schema::Encoding;
+use parquet2::encoding::Encoding;
 use parquet2::{metadata::ColumnDescriptor, page::CompressedDataPage, write::WriteOptions};
 
 use super::super::{levels, utils};

diff --git a/src/io/parquet/write/boolean/basic.rs b/src/io/parquet/write/boolean/basic.rs
@@ -1,8 +1,7 @@
 use parquet2::{
-    encoding::hybrid_rle::bitpacked_encode,
+    encoding::{hybrid_rle::bitpacked_encode, Encoding},
     metadata::ColumnDescriptor,
     page::CompressedDataPage,
-    schema::Encoding,
     statistics::{serialize_statistics, BooleanStatistics, ParquetStatistics, Statistics},
     write::WriteOptions,
 };

diff --git a/src/io/parquet/write/boolean/nested.rs b/src/io/parquet/write/boolean/nested.rs
@@ -1,4 +1,4 @@
-use parquet2::schema::Encoding;
+use parquet2::encoding::Encoding;
 use parquet2::{metadata::ColumnDescriptor, page::CompressedDataPage, write::WriteOptions};
 
 use super::super::{levels, utils};

diff --git a/src/io/parquet/write/dictionary.rs b/src/io/parquet/write/dictionary.rs
@@ -1,6 +1,5 @@
-use parquet2::encoding::hybrid_rle::encode_u32;
+use parquet2::encoding::{hybrid_rle::encode_u32, Encoding};
 use parquet2::page::{CompressedDictPage, CompressedPage};
-use parquet2::schema::Encoding;
 use parquet2::write::DynIter;
 use parquet2::{metadata::ColumnDescriptor, write::WriteOptions};
 

diff --git a/src/io/parquet/write/fixed_len_bytes.rs b/src/io/parquet/write/fixed_len_bytes.rs
@@ -1,6 +1,6 @@
 use parquet2::{
-    compression::create_codec, metadata::ColumnDescriptor, page::CompressedDataPage,
-    schema::Encoding, write::WriteOptions,
+    compression::create_codec, encoding::Encoding, metadata::ColumnDescriptor,
+    page::CompressedDataPage, write::WriteOptions,
 };
 
 use super::utils;