Repository: jorgecarleitao/arrow2
Pull request #1312 (2022-12-01 14:02:42)

Fixed failing to accept dictionary full of nulls
This fixes https://github.com/pola-rs/polars/issues/5655

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1312?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
Base: **83.14**% // Head: **83.15**% // Increases project coverage by **`+0.00%`** :tada:
> Coverage data is based on head [(`53f265d`)](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1312?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) compared to base [(`6d2eccd`)](https://codecov.io/gh/jorgecarleitao/arrow2/commit/6d2eccd78a167b39e8a5a5e8a3f053422104bee9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).
> Patch coverage: 100.00% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main    #1312   +/-   ##
=======================================
  Coverage   83.14%   83.15%           
=======================================
  Files         369      369           
  Lines       40245    40247    +2     
=======================================
+ Hits        33461    33466    +5     
+ Misses       6784     6781    -3     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1312?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/array/dictionary/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2RpY3Rpb25hcnkvbW9kLnJz) | `91.76% <100.00%> (+0.09%)` | :arrow_up: |
| [src/array/utf8/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1312/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3V0ZjgvbW9kLnJz) | `85.67% <0.00%> (+0.91%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1312?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).


Patch:
diff --git a/Cargo.toml b/Cargo.toml
index 7cf1ccb238..f0f2714a0f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -77,9 +77,6 @@ futures = { version = "0.3", optional = true }
 # to read IPC as a stream
 async-stream = { version = "0.3.2", optional = true }
 
-# parquet support
-parquet2 = { version = "0.16", optional = true, default_features = false, features = ["async"] }
-
 # avro support
 avro-schema = { version = "0.3", optional = true }
 
@@ -108,6 +105,13 @@ getrandom = { version = "0.2", features = ["js"] }
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 ahash = { version = "0.8", features=["runtime-rng"] }
 
+# parquet support
+[dependencies.parquet2]
+version = "0.17"
+optional = true
+default_features = false
+features = ["async"]
+
 [dev-dependencies]
 criterion = "0.3"
 flate2 = "1"
diff --git a/src/array/dictionary/mod.rs b/src/array/dictionary/mod.rs
index 8b62c91dce..63cdfc7fa7 100644
--- a/src/array/dictionary/mod.rs
+++ b/src/array/dictionary/mod.rs
@@ -119,7 +119,9 @@ impl<K: DictionaryKey> DictionaryArray<K> {
     ) -> Result<Self, Error> {
         check_data_type(K::KEY_TYPE, &data_type, values.data_type())?;
 
-        check_indexes(keys.values(), values.len())?;
+        if keys.null_count() != keys.len() {
+            check_indexes(keys.values(), values.len())?;
+        }
 
         Ok(Self {
             data_type,
diff --git a/src/io/ipc/write/common_async.rs b/src/io/ipc/write/common_async.rs
index 21c04603d8..43c7ff73d6 100644
--- a/src/io/ipc/write/common_async.rs
+++ b/src/io/ipc/write/common_async.rs
@@ -48,7 +48,6 @@ pub async fn write_continuation<W: AsyncWrite + Unpin + Send>(
 ) -> Result<usize> {
     writer.write_all(&CONTINUATION_MARKER).await?;
     writer.write_all(&total_len.to_le_bytes()[..]).await?;
-    writer.flush().await?;
     Ok(8)
 }
 
@@ -66,6 +65,5 @@ async fn write_body_buffers<W: AsyncWrite + Unpin + Send>(
         writer.write_all(&vec![0u8; pad_len][..]).await?;
     }
 
-    writer.flush().await?;
     Ok(total_len)
 }
diff --git a/src/io/ipc/write/common_sync.rs b/src/io/ipc/write/common_sync.rs
index 9f281e2206..6d3966948d 100644
--- a/src/io/ipc/write/common_sync.rs
+++ b/src/io/ipc/write/common_sync.rs
@@ -47,7 +47,6 @@ fn write_body_buffers<W: Write>(mut writer: W, data: &[u8]) -> Result<usize> {
         writer.write_all(&vec![0u8; pad_len][..])?;
     }
 
-    writer.flush()?;
     Ok(total_len)
 }
 
@@ -56,6 +55,5 @@ fn write_body_buffers<W: Write>(mut writer: W, data: &[u8]) -> Result<usize> {
 pub fn write_continuation<W: Write>(writer: &mut W, total_len: i32) -> Result<usize> {
     writer.write_all(&CONTINUATION_MARKER)?;
     writer.write_all(&total_len.to_le_bytes()[..])?;
-    writer.flush()?;
     Ok(8)
 }
diff --git a/src/io/parquet/read/mod.rs b/src/io/parquet/read/mod.rs
index be0ddc9ac5..b8031a3758 100644
--- a/src/io/parquet/read/mod.rs
+++ b/src/io/parquet/read/mod.rs
@@ -22,8 +22,7 @@ pub use parquet2::{
         decompress, get_column_iterator, get_page_stream,
         read_columns_indexes as _read_columns_indexes, read_metadata as _read_metadata,
         read_metadata_async as _read_metadata_async, read_pages_locations, BasicDecompressor,
-        ColumnChunkIter, Decompressor, MutStreamingIterator, PageFilter, PageReader,
-        ReadColumnIterator, State,
+        Decompressor, MutStreamingIterator, PageFilter, PageReader, ReadColumnIterator, State,
     },
     schema::types::{
         GroupLogicalType, ParquetType, PhysicalType, PrimitiveConvertedType, PrimitiveLogicalType,
diff --git a/src/io/parquet/write/dictionary.rs b/src/io/parquet/write/dictionary.rs
index 2abffd98a4..9f20d4692d 100644
--- a/src/io/parquet/write/dictionary.rs
+++ b/src/io/parquet/write/dictionary.rs
@@ -1,6 +1,6 @@
 use parquet2::{
     encoding::{hybrid_rle::encode_u32, Encoding},
-    page::{DictPage, EncodedPage},
+    page::{DictPage, Page},
     schema::types::PrimitiveType,
     statistics::{serialize_statistics, ParquetStatistics},
     write::DynIter,
@@ -106,7 +106,7 @@ fn serialize_keys<K: DictionaryKey>(
     nested: &[Nested],
     statistics: ParquetStatistics,
     options: WriteOptions,
-) -> Result<EncodedPage> {
+) -> Result<Page> {
     let mut buffer = vec![];
 
     // parquet only accepts a single validity - we "&" the validities into a single one
@@ -142,7 +142,7 @@ fn serialize_keys<K: DictionaryKey>(
         options,
         Encoding::RleDictionary,
     )
-    .map(EncodedPage::Data)
+    .map(Page::Data)
 }
 
 macro_rules! dyn_prim {
@@ -162,7 +162,7 @@ pub fn array_to_pages<K: DictionaryKey>(
     nested: &[Nested],
     options: WriteOptions,
     encoding: Encoding,
-) -> Result<DynIter<'static, Result<EncodedPage>>> {
+) -> Result<DynIter<'static, Result<Page>>> {
     match encoding {
         Encoding::PlainDictionary | Encoding::RleDictionary => {
             // write DictPage
@@ -230,7 +230,7 @@ pub fn array_to_pages<K: DictionaryKey>(
                     )))
                 }
             };
-            let dict_page = EncodedPage::Dict(dict_page);
+            let dict_page = Page::Dict(dict_page);
 
             // write DataPage pointing to DictPage
             let data_page = serialize_keys(array, type_, nested, statistics, options)?;
diff --git a/src/io/parquet/write/mod.rs b/src/io/parquet/write/mod.rs
index f48c1a67ea..84b4f1cab5 100644
--- a/src/io/parquet/write/mod.rs
+++ b/src/io/parquet/write/mod.rs
@@ -38,7 +38,7 @@ pub use parquet2::{
     encoding::Encoding,
     fallible_streaming_iterator,
     metadata::{Descriptor, FileMetaData, KeyValue, SchemaDescriptor, ThriftFileMetaData},
-    page::{CompressedDataPage, CompressedPage, EncodedPage},
+    page::{CompressedDataPage, CompressedPage, Page},
     schema::types::{FieldInfo, ParquetType, PhysicalType as ParquetPhysicalType},
     write::{
         compress, write_metadata_sidecar, Compressor, DynIter, DynStreamingIterator, RowGroupIter,
@@ -130,7 +130,7 @@ pub fn can_encode(data_type: &DataType, encoding: Encoding) -> bool {
     )
 }
 
-/// Returns an iterator of [`EncodedPage`].
+/// Returns an iterator of [`Page`].
 #[allow(clippy::needless_collect)]
 pub fn array_to_pages(
     array: &dyn Array,
@@ -138,7 +138,7 @@ pub fn array_to_pages(
     nested: &[Nested],
     options: WriteOptions,
     encoding: Encoding,
-) -> Result<DynIter<'static, Result<EncodedPage>>> {
+) -> Result<DynIter<'static, Result<Page>>> {
     // maximum page size is 2^31 e.g. i32::MAX
     // we split at 2^31 - 2^25 to err on the safe side
     // we also check for an array.len > 3 to prevent infinite recursion
@@ -175,7 +175,7 @@ pub fn array_to_pages(
                     ((array_byte_size as f64) / ((array.len() + 1) as f64)) as usize;
                 let rows_per_page = (page_size / (bytes_per_row + 1)).max(1);
 
-                let vs: Vec<Result<EncodedPage>> = (0..array.len())
+                let vs: Vec<Result<Page>> = (0..array.len())
                     .step_by(rows_per_page)
                     .map(|offset| {
                         let length = if offset + rows_per_page > array.len() {
@@ -202,7 +202,7 @@ pub fn array_to_page(
     nested: &[Nested],
     options: WriteOptions,
     encoding: Encoding,
-) -> Result<EncodedPage> {
+) -> Result<Page> {
     if nested.len() == 1 {
         // special case where validity == def levels
         return array_to_page_simple(array, type_, options, encoding);
@@ -216,7 +216,7 @@ pub fn array_to_page_simple(
     type_: ParquetPrimitiveType,
     options: WriteOptions,
     encoding: Encoding,
-) -> Result<EncodedPage> {
+) -> Result<Page> {
     let data_type = array.data_type();
     if !can_encode(data_type, encoding) {
         return Err(Error::InvalidArgumentError(format!(
@@ -439,7 +439,7 @@ pub fn array_to_page_simple(
             other
         ))),
     }
-    .map(EncodedPage::Data)
+    .map(Page::Data)
 }
 
 fn array_to_page_nested(
@@ -448,7 +448,7 @@ fn array_to_page_nested(
     nested: &[Nested],
     options: WriteOptions,
     _encoding: Encoding,
-) -> Result<EncodedPage> {
+) -> Result<Page> {
     use DataType::*;
     match array.data_type().to_logical_type() {
         Null => {
@@ -520,7 +520,7 @@ fn array_to_page_nested(
             other
         ))),
     }
-    .map(EncodedPage::Data)
+    .map(Page::Data)
 }
 
 fn transverse_recursive<T, F: Fn(&DataType) -> T + Clone>(
diff --git a/src/io/parquet/write/pages.rs b/src/io/parquet/write/pages.rs
index 5e0c80112d..caaa92866b 100644
--- a/src/io/parquet/write/pages.rs
+++ b/src/io/parquet/write/pages.rs
@@ -1,5 +1,5 @@
 use parquet2::schema::types::{ParquetType, PrimitiveType as ParquetPrimitiveType};
-use parquet2::{page::EncodedPage, write::DynIter};
+use parquet2::{page::Page, write::DynIter};
 
 use crate::array::{ListArray, Offset, StructArray};
 use crate::bitmap::Bitmap;
@@ -193,13 +193,13 @@ fn to_parquet_leafs_recursive(type_: ParquetType, leafs: &mut Vec<ParquetPrimiti
     }
 }
 
-/// Returns a vector of iterators of [`EncodedPage`], one per leaf column in the array
+/// Returns a vector of iterators of [`Page`], one per leaf column in the array
 pub fn array_to_columns<A: AsRef<dyn Array> + Send + Sync>(
     array: A,
     type_: ParquetType,
     options: WriteOptions,
     encoding: &[Encoding],
-) -> Result<Vec<DynIter<'static, Result<EncodedPage>>>> {
+) -> Result<Vec<DynIter<'static, Result<Page>>>> {
     let array = array.as_ref();
     let nested = to_nested(array, &type_)?;
 


Test patch:
diff --git a/tests/it/array/dictionary/mod.rs b/tests/it/array/dictionary/mod.rs
index 84c16b49cd..b139c94c53 100644
--- a/tests/it/array/dictionary/mod.rs
+++ b/tests/it/array/dictionary/mod.rs
@@ -40,6 +40,20 @@ fn try_new_incorrect_key() {
     assert!(r);
 }
 
+#[test]
+fn try_new_nulls() {
+    let key: Option<u32> = None;
+    let keys = PrimitiveArray::from_iter([key]);
+    let value: &[&str] = &[];
+    let values = Utf8Array::<i32>::from_slice(value);
+
+    let data_type =
+        DataType::Dictionary(u32::KEY_TYPE, Box::new(values.data_type().clone()), false);
+    let r = DictionaryArray::try_new(data_type, keys, values.boxed()).is_ok();
+
+    assert!(r);
+}
+
 #[test]
 fn try_new_incorrect_dt() {
     let values = Utf8Array::<i32>::from_slice(["a", "aa"]);
diff --git a/tests/it/io/parquet/read_indexes.rs b/tests/it/io/parquet/read_indexes.rs
index b17803b16f..4e41bb2baf 100644
--- a/tests/it/io/parquet/read_indexes.rs
+++ b/tests/it/io/parquet/read_indexes.rs
@@ -6,10 +6,7 @@ use arrow2::io::parquet::read::indexes;
 use arrow2::{array::*, datatypes::*, error::Result, io::parquet::read::*, io::parquet::write::*};
 
 /// Returns 2 sets of pages with different the same number of rows distributed un-evenly
-fn pages(
-    arrays: &[&dyn Array],
-    encoding: Encoding,
-) -> Result<(Vec<EncodedPage>, Vec<EncodedPage>, Schema)> {
+fn pages(arrays: &[&dyn Array], encoding: Encoding) -> Result<(Vec<Page>, Vec<Page>, Schema)> {
     // create pages with different number of rows
     let array11 = PrimitiveArray::<i64>::from_slice([1, 2, 3, 4]);
     let array12 = PrimitiveArray::<i64>::from_slice([5]);
@@ -73,7 +70,7 @@ fn pages(
 
 /// Tests reading pages while skipping indexes
 fn read_with_indexes(
-    (pages1, pages2, schema): (Vec<EncodedPage>, Vec<EncodedPage>, Schema),
+    (pages1, pages2, schema): (Vec<Page>, Vec<Page>, Schema),
     expected: Box<dyn Array>,
 ) -> Result<()> {
     let options = WriteOptions {
@@ -83,7 +80,7 @@ fn read_with_indexes(
         data_pagesize_limit: None,
     };
 
-    let to_compressed = |pages: Vec<EncodedPage>| {
+    let to_compressed = |pages: Vec<Page>| {
         let encoded_pages = DynIter::new(pages.into_iter().map(Ok));
         let compressed_pages =
             Compressor::new(encoded_pages, options.compression, vec![]).map_err(Error::from);