Migrate to latest parquet2

Cargo.toml

@@ -53,7 +53,7 @@ hex = { version = "^0.4", optional = true }
 
 # for IPC compression
 lz4 = { version = "1.23.1", optional = true }
-zstd = { version = "0.10", optional = true }
+zstd = { version = "0.11", optional = true }
 
 rand = { version = "0.8", optional = true }
 
@@ -68,7 +68,8 @@ futures = { version = "0.3", optional = true }
 ahash = { version = "0.7", optional = true }
 
 # parquet support
-parquet2 = { version = "0.10", optional = true, default_features = false, features = ["stream"] }
+#parquet2 = { version = "0.10", optional = true, default_features = false, features = ["stream"] }
+parquet2 = { git = "https://github.com/jorgecarleitao/parquet2", branch = "write_indexes", optional = true, default_features = false, features = ["stream"] }
 
 # avro support
 avro-schema = { version = "0.2", optional = true }

arrow-parquet-integration-testing/src/main.rs

@@ -196,8 +196,7 @@ fn main() -> Result<()> {
 
     writer.start()?;
     for group in row_groups {
-        let (group, len) = group?;
-        writer.write(group, len)?;
+        writer.write(group?)?;
     }
     let _ = writer.end(None)?;
 

benches/write_parquet.rs

@@ -34,8 +34,7 @@ fn write(array: &dyn Array, encoding: Encoding) -> Result<()> {
 
     writer.start()?;
     for group in row_groups {
-        let (group, len) = group?;
-        writer.write(group, len)?;
+        writer.write(group?)?;
     }
     let _ = writer.end(None)?;
     Ok(())

examples/parquet_write.rs

@@ -30,8 +30,7 @@ fn write_batch(path: &str, schema: Schema, columns: Chunk<Arc<dyn Array>>) -> Re
 
     writer.start()?;
     for group in row_groups {
-        let (group, len) = group?;
-        writer.write(group, len)?;
+        writer.write(group?)?;
     }
     let _size = writer.end(None)?;
     Ok(())

examples/parquet_write_parallel/src/main.rs

@@ -99,8 +99,7 @@ fn parallel_write(path: &str, schema: &Schema, batches: &[Chunk]) -> Result<()>
     // Write the file.
     writer.start()?;
     for group in row_groups {
-        let (group, len) = group?;
-        writer.write(group, len)?;
+        writer.write(group?)?;
     }
     let _size = writer.end(None)?;
 

src/doc/lib.md

@@ -62,8 +62,7 @@ fn main() -> Result<()> {
     // Write the file.
     writer.start()?;
     for group in row_groups {
-        let (group, len) = group?;
-        writer.write(group, len)?;
+        writer.write(group?)?;
     }
     let _ = writer.end(None)?;
     Ok(())

src/error.rs

@@ -52,6 +52,12 @@ impl From<std::str::Utf8Error> for ArrowError {
     }
 }
 
+impl From<std::string::FromUtf8Error> for ArrowError {
+    fn from(error: std::string::FromUtf8Error) -> Self {
+        ArrowError::External("".to_string(), Box::new(error))
+    }
+}
+
 impl From<simdutf8::basic::Utf8Error> for ArrowError {
     fn from(error: simdutf8::basic::Utf8Error) -> Self {
         ArrowError::External("".to_string(), Box::new(error))

src/io/parquet/read/deserialize/binary/basic.rs

@@ -180,7 +180,7 @@ impl<'a, O: Offset> utils::Decoder<'a> for BinaryDecoder<O> {
 
     fn build_state(&self, page: &'a DataPage) -> Result<Self::State> {
         let is_optional =
-            page.descriptor().type_().get_basic_info().repetition() == &Repetition::Optional;
+            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
 
         match (page.encoding(), page.dictionary_page(), is_optional) {
             (Encoding::PlainDictionary | Encoding::RleDictionary, Some(dict), false) => {

src/io/parquet/read/deserialize/binary/nested.rs

@@ -42,7 +42,7 @@ impl<'a, O: Offset> utils::Decoder<'a> for BinaryDecoder<O> {
 
     fn build_state(&self, page: &'a DataPage) -> Result<Self::State> {
         let is_optional =
-            page.descriptor().type_().get_basic_info().repetition() == &Repetition::Optional;
+            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
 
         match (page.encoding(), page.dictionary_page(), is_optional) {
             (Encoding::Plain, None, true) => {

src/io/parquet/read/deserialize/boolean/basic.rs

@@ -89,7 +89,7 @@ impl<'a> Decoder<'a> for BooleanDecoder {
 
     fn build_state(&self, page: &'a DataPage) -> Result<Self::State> {
         let is_optional =
-            page.descriptor().type_().get_basic_info().repetition() == &Repetition::Optional;
+            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
 
         match (page.encoding(), is_optional) {
             (Encoding::Plain, true) => Ok(State::Optional(Optional::new(page))),

src/io/parquet/read/deserialize/boolean/nested.rs

@@ -65,7 +65,7 @@ impl<'a> Decoder<'a> for BooleanDecoder {
 
     fn build_state(&self, page: &'a DataPage) -> Result<Self::State> {
         let is_optional =
-            page.descriptor().type_().get_basic_info().repetition() == &Repetition::Optional;
+            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
 
         match (page.encoding(), is_optional) {
             (Encoding::Plain, true) => {

src/io/parquet/read/deserialize/dictionary.rs

@@ -131,7 +131,7 @@ where
 
     fn build_state(&self, page: &'a DataPage) -> Result<Self::State> {
         let is_optional =
-            page.descriptor().type_().get_basic_info().repetition() == &Repetition::Optional;
+            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
 
         match (page.encoding(), is_optional) {
             (Encoding::PlainDictionary | Encoding::RleDictionary, false) => {

src/io/parquet/read/deserialize/fixed_size_binary/basic.rs

@@ -121,7 +121,7 @@ impl<'a> Decoder<'a> for BinaryDecoder {
 
     fn build_state(&self, page: &'a DataPage) -> Result<Self::State> {
         let is_optional =
-            page.descriptor().type_().get_basic_info().repetition() == &Repetition::Optional;
+            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
 
         match (page.encoding(), page.dictionary_page(), is_optional) {
             (Encoding::Plain, None, true) => Ok(State::Optional(Optional::new(page, self.size))),

src/io/parquet/read/deserialize/mod.rs

@@ -17,6 +17,7 @@ use crate::{
 };
 
 use self::nested_utils::{InitNested, NestedArrayIter, NestedState};
+use parquet2::schema::types::PrimitiveType;
 use simple::page_iter_to_arrays;
 
 use super::*;
@@ -27,7 +28,7 @@ pub fn get_page_iterator<R: Read + Seek>(
     reader: R,
     pages_filter: Option<PageFilter>,
     buffer: Vec<u8>,
-) -> Result<PageIterator<R>> {
+) -> Result<PageReader<R>> {
     Ok(_get_page_iterator(
         column_metadata,
         reader,
@@ -76,7 +77,7 @@ fn create_list(
 
 fn columns_to_iter_recursive<'a, I: 'a>(
     mut columns: Vec<I>,
-    mut types: Vec<&ParquetType>,
+    mut types: Vec<&PrimitiveType>,
     field: Field,
     mut init: Vec<InitNested>,
     chunk_size: usize,
@@ -238,7 +239,7 @@ fn field_to_init(field: &Field) -> Vec<InitNested> {
 /// The arrays are guaranteed to be at most of size `chunk_size` and data type `field.data_type`.
 pub fn column_iter_to_arrays<'a, I: 'a>(
     columns: Vec<I>,
-    types: Vec<&ParquetType>,
+    types: Vec<&PrimitiveType>,
     field: Field,
     chunk_size: usize,
 ) -> Result<ArrayIter<'a>>

src/io/parquet/read/deserialize/nested_utils.rs

@@ -283,8 +283,8 @@ impl<'a> NestedPage<'a> {
     pub fn new(page: &'a DataPage) -> Self {
         let (rep_levels, def_levels, _) = split_buffer(page);
 
-        let max_rep_level = page.descriptor().max_rep_level();
-        let max_def_level = page.descriptor().max_def_level();
+        let max_rep_level = page.descriptor.max_rep_level;
+        let max_def_level = page.descriptor.max_def_level;
 
         let reps =
             HybridRleDecoder::new(rep_levels, get_bit_width(max_rep_level), page.num_values());
@@ -451,7 +451,7 @@ impl<'a> Optional<'a> {
     pub fn new(page: &'a DataPage) -> Self {
         let (_, def_levels, _) = split_buffer(page);
 
-        let max_def = page.descriptor().max_def_level();
+        let max_def = page.descriptor.max_def_level;
 
         Self {
             definition_levels: HybridRleDecoder::new(

src/io/parquet/read/deserialize/primitive/basic.rs

@@ -137,7 +137,7 @@ where
 
     fn build_state(&self, page: &'a DataPage) -> Result<Self::State> {
         let is_optional =
-            page.descriptor().type_().get_basic_info().repetition() == &Repetition::Optional;
+            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
 
         match (page.encoding(), page.dictionary_page(), is_optional) {
             (Encoding::PlainDictionary | Encoding::RleDictionary, Some(dict), false) => {

src/io/parquet/read/deserialize/primitive/nested.rs

@@ -81,7 +81,7 @@ where
 
     fn build_state(&self, page: &'a DataPage) -> Result<Self::State> {
         let is_optional =
-            page.descriptor().type_().get_basic_info().repetition() == &Repetition::Optional;
+            page.descriptor.primitive_type.field_info.repetition == Repetition::Optional;
 
         match (page.encoding(), page.dictionary_page(), is_optional) {
             (Encoding::PlainDictionary | Encoding::RleDictionary, Some(dict), false) => {

src/io/parquet/read/deserialize/simple.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use parquet2::{
     schema::types::{
-        LogicalType, ParquetType, PhysicalType, TimeUnit as ParquetTimeUnit, TimestampType,
+        LogicalType, PhysicalType, PrimitiveType, TimeUnit as ParquetTimeUnit, TimestampType,
     },
     types::int96_to_i64_ns,
 };
@@ -60,24 +60,14 @@ where
 /// of [`DataType`] `data_type` and `chunk_size`.
 pub fn page_iter_to_arrays<'a, I: 'a + DataPages>(
     pages: I,
-    type_: &ParquetType,
+    type_: &PrimitiveType,
     data_type: DataType,
     chunk_size: usize,
 ) -> Result<ArrayIter<'a>> {
     use DataType::*;
 
-    let (physical_type, logical_type) = if let ParquetType::PrimitiveType {
-        physical_type,
-        logical_type,
-        ..
-    } = type_
-    {
-        (physical_type, logical_type)
-    } else {
-        return Err(ArrowError::InvalidArgumentError(
-            "page_iter_to_arrays can only be called with a parquet primitive type".into(),
-        ));
-    };
+    let physical_type = &type_.physical_type;
+    let logical_type = &type_.logical_type;
 
     Ok(match data_type.to_logical_type() {
         Null => null::iter_to_arrays(pages, data_type, chunk_size),
@@ -276,7 +266,7 @@ fn timestamp<'a, I: 'a + DataPages>(
     Ok(match (unit, time_unit) {
         (ParquetTimeUnit::MILLIS(_), TimeUnit::Second) => dyn_iter(op(iter, |x| x / 1_000)),
         (ParquetTimeUnit::MICROS(_), TimeUnit::Second) => dyn_iter(op(iter, |x| x / 1_000_000)),
-        (ParquetTimeUnit::NANOS(_), TimeUnit::Second) => dyn_iter(op(iter, |x| x * 1_000_000_000)),
+        (ParquetTimeUnit::NANOS(_), TimeUnit::Second) => dyn_iter(op(iter, |x| x / 1_000_000_000)),
 
         (ParquetTimeUnit::MILLIS(_), TimeUnit::Millisecond) => dyn_iter(iden(iter)),
         (ParquetTimeUnit::MICROS(_), TimeUnit::Millisecond) => dyn_iter(op(iter, |x| x / 1_000)),
@@ -348,7 +338,7 @@ fn timestamp_dict<'a, K: DictionaryKey, I: 'a + DataPages>(
                 pages,
                 data_type,
                 chunk_size,
-                |x: i64| x * 1_000_000_000,
+                |x: i64| x / 1_000_000_000,
             ))
         }
 

src/io/parquet/read/deserialize/utils.rs

@@ -55,7 +55,7 @@ pub fn not_implemented(
 
 #[inline]
 pub fn split_buffer(page: &DataPage) -> (&[u8], &[u8], &[u8]) {
-    _split_buffer(page, page.descriptor())
+    _split_buffer(page)
 }
 
 /// A private trait representing structs that can receive elements.

src/io/parquet/read/indexes/binary.rs

@@ -0,0 +1,43 @@
+use parquet2::indexes::PageIndex;
+
+use crate::{
+    array::{Array, BinaryArray, PrimitiveArray, Utf8Array},
+    datatypes::{DataType, PhysicalType},
+    error::ArrowError,
+    trusted_len::TrustedLen,
+};
+
+use super::ColumnIndex;
+
+pub fn deserialize(
+    indexes: &[PageIndex<Vec<u8>>],
+    data_type: &DataType,
+) -> Result<ColumnIndex, ArrowError> {
+    Ok(ColumnIndex {
+        min: deserialize_binary_iter(indexes.iter().map(|index| index.min.as_ref()), data_type)?,
+        max: deserialize_binary_iter(indexes.iter().map(|index| index.max.as_ref()), data_type)?,
+        null_count: PrimitiveArray::from_trusted_len_iter(
+            indexes
+                .iter()
+                .map(|index| index.null_count.map(|x| x as u64)),
+        ),
+    })
+}
+
+fn deserialize_binary_iter<'a, I: TrustedLen<Item = Option<&'a Vec<u8>>>>(
+    iter: I,
+    data_type: &DataType,
+) -> Result<Box<dyn Array>, ArrowError> {
+    match data_type.to_physical_type() {
+        PhysicalType::LargeBinary => Ok(Box::new(BinaryArray::<i64>::from_iter(iter))),
+        PhysicalType::Utf8 => {
+            let iter = iter.map(|x| x.map(|x| std::str::from_utf8(x)).transpose());
+            Ok(Box::new(Utf8Array::<i32>::try_from_trusted_len_iter(iter)?))
+        }
+        PhysicalType::LargeUtf8 => {
+            let iter = iter.map(|x| x.map(|x| std::str::from_utf8(x)).transpose());
+            Ok(Box::new(Utf8Array::<i64>::try_from_trusted_len_iter(iter)?))
+        }
+        _ => Ok(Box::new(BinaryArray::<i32>::from_iter(iter))),
+    }
+}

src/io/parquet/read/indexes/boolean.rs

@@ -0,0 +1,21 @@
+use parquet2::indexes::PageIndex;
+
+use crate::array::{BooleanArray, PrimitiveArray};
+
+use super::ColumnIndex;
+
+pub fn deserialize(indexes: &[PageIndex<bool>]) -> ColumnIndex {
+    ColumnIndex {
+        min: Box::new(BooleanArray::from_trusted_len_iter(
+            indexes.iter().map(|index| index.min),
+        )),
+        max: Box::new(BooleanArray::from_trusted_len_iter(
+            indexes.iter().map(|index| index.max),
+        )),
+        null_count: PrimitiveArray::from_trusted_len_iter(
+            indexes
+                .iter()
+                .map(|index| index.null_count.map(|x| x as u64)),
+        ),
+    }
+}

src/io/parquet/read/indexes/fixed_len_binary.rs

@@ -0,0 +1,58 @@
+use parquet2::indexes::PageIndex;
+
+use crate::{
+    array::{Array, FixedSizeBinaryArray, MutableFixedSizeBinaryArray, PrimitiveArray},
+    datatypes::{DataType, PhysicalType, PrimitiveType},
+    trusted_len::TrustedLen,
+};
+
+use super::ColumnIndex;
+
+pub fn deserialize(indexes: &[PageIndex<Vec<u8>>], data_type: DataType) -> ColumnIndex {
+    ColumnIndex {
+        min: deserialize_binary_iter(
+            indexes.iter().map(|index| index.min.as_ref()),
+            data_type.clone(),
+        ),
+        max: deserialize_binary_iter(indexes.iter().map(|index| index.max.as_ref()), data_type),
+        null_count: PrimitiveArray::from_trusted_len_iter(
+            indexes
+                .iter()
+                .map(|index| index.null_count.map(|x| x as u64)),
+        ),
+    }
+}
+
+fn deserialize_binary_iter<'a, I: TrustedLen<Item = Option<&'a Vec<u8>>>>(
+    iter: I,
+    data_type: DataType,
+) -> Box<dyn Array> {
+    match data_type.to_physical_type() {
+        PhysicalType::Primitive(PrimitiveType::Int128) => {
+            Box::new(PrimitiveArray::from_trusted_len_iter(iter.map(|v| {
+                v.map(|x| {
+                    // Copy the fixed-size byte value to the start of a 16 byte stack
+                    // allocated buffer, then use an arithmetic right shift to fill in
+                    // MSBs, which accounts for leading 1's in negative (two's complement)
+                    // values.
+                    let n = x.len();
+                    let mut bytes = [0u8; 16];
+                    bytes[..n].copy_from_slice(x);
+                    i128::from_be_bytes(bytes) >> (8 * (16 - n))
+                })
+            })))
+        }
+        _ => {
+            let mut a = MutableFixedSizeBinaryArray::from_data(
+                data_type,
+                Vec::with_capacity(iter.size_hint().0),
+                None,
+            );
+            for item in iter {
+                a.push(item);
+            }
+            let a: FixedSizeBinaryArray = a.into();
+            Box::new(a)
+        }
+    }
+}