Repository: jorgecarleitao/arrow2
Pull request #1113 (2022-06-27 22:31:20)

Fixed panic in writing dictionaries to parquet
Also added tests covering them (and others)

codecov[bot]: # [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) Report
> Merging [#1113](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (b21981b) into [main](https://codecov.io/gh/jorgecarleitao/arrow2/commit/b942a84d0e8d3980f5e754801371c5d2ee65551a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) (b942a84) will **increase** coverage by `0.41%`.
> The diff coverage is `n/a`.

```diff
@@            Coverage Diff             @@
##             main    #1113      +/-   ##
==========================================
+ Coverage   81.33%   81.74%   +0.41%     
==========================================
  Files         367      367              
  Lines       35458    35830     +372     
==========================================
+ Hits        28840    29290     +450     
+ Misses       6618     6540      -78     
```


| [Impacted Files](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | Coverage Δ | |
|---|---|---|
| [src/io/parquet/read/deserialize/boolean/basic.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL3BhcnF1ZXQvcmVhZC9kZXNlcmlhbGl6ZS9ib29sZWFuL2Jhc2ljLnJz) | `93.41% <0.00%> (-1.08%)` | :arrow_down: |
| [...rc/io/parquet/read/deserialize/primitive/nested.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL3BhcnF1ZXQvcmVhZC9kZXNlcmlhbGl6ZS9wcmltaXRpdmUvbmVzdGVkLnJz) | `87.96% <0.00%> (-0.27%)` | :arrow_down: |
| [src/io/parquet/read/deserialize/nested\_utils.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL3BhcnF1ZXQvcmVhZC9kZXNlcmlhbGl6ZS9uZXN0ZWRfdXRpbHMucnM=) | `83.71% <0.00%> (-0.23%)` | :arrow_down: |
| [src/io/ipc/read/reader.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL2lwYy9yZWFkL3JlYWRlci5ycw==) | `96.02% <0.00%> (ø)` | |
| [src/array/utf8/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L3V0ZjgvbW9kLnJz) | `91.07% <0.00%> (+0.27%)` | :arrow_up: |
| [src/io/parquet/write/dictionary.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL3BhcnF1ZXQvd3JpdGUvZGljdGlvbmFyeS5ycw==) | `81.94% <0.00%> (+0.69%)` | :arrow_up: |
| [src/io/parquet/read/deserialize/boolean/nested.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL3BhcnF1ZXQvcmVhZC9kZXNlcmlhbGl6ZS9ib29sZWFuL25lc3RlZC5ycw==) | `68.10% <0.00%> (+1.04%)` | :arrow_up: |
| [src/array/binary/mod.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2FycmF5L2JpbmFyeS9tb2QucnM=) | `90.12% <0.00%> (+1.23%)` | :arrow_up: |
| [src/io/parquet/read/deserialize/binary/basic.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL3BhcnF1ZXQvcmVhZC9kZXNlcmlhbGl6ZS9iaW5hcnkvYmFzaWMucnM=) | `76.47% <0.00%> (+1.27%)` | :arrow_up: |
| [...arquet/read/deserialize/fixed\_size\_binary/basic.rs](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao#diff-c3JjL2lvL3BhcnF1ZXQvcmVhZC9kZXNlcmlhbGl6ZS9maXhlZF9zaXplX2JpbmFyeS9iYXNpYy5ycw==) | `97.20% <0.00%> (+1.37%)` | :arrow_up: |
| ... and [8 more](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Last update [b942a84...b21981b](https://codecov.io/gh/jorgecarleitao/arrow2/pull/1113?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Jorge+Leitao).


Patch:
diff --git a/src/io/parquet/read/deserialize/simple.rs b/src/io/parquet/read/deserialize/simple.rs
index aaab8aee9f..20dba38065 100644
--- a/src/io/parquet/read/deserialize/simple.rs
+++ b/src/io/parquet/read/deserialize/simple.rs
@@ -467,6 +467,12 @@ fn dict_read<'a, K: DictionaryKey, I: 'a + DataPages>(
             chunk_size,
             |x: i32| x as u32,
         )),
+        UInt64 => dyn_iter(primitive::DictIter::<K, _, _, _, _>::new(
+            iter,
+            data_type,
+            chunk_size,
+            |x: i64| x as u64,
+        )),
         Int8 => dyn_iter(primitive::DictIter::<K, _, _, _, _>::new(
             iter,
             data_type,
diff --git a/src/io/parquet/write/dictionary.rs b/src/io/parquet/write/dictionary.rs
index 65f13a4904..05c56207bf 100644
--- a/src/io/parquet/write/dictionary.rs
+++ b/src/io/parquet/write/dictionary.rs
@@ -149,7 +149,7 @@ pub fn array_to_pages<K: DictionaryKey>(
                 DataType::UInt8 => dyn_prim!(u8, i32, array, options, type_),
                 DataType::UInt16 => dyn_prim!(u16, i32, array, options, type_),
                 DataType::UInt32 => dyn_prim!(u32, i32, array, options, type_),
-                DataType::UInt64 => dyn_prim!(i64, i64, array, options, type_),
+                DataType::UInt64 => dyn_prim!(u64, i64, array, options, type_),
                 DataType::Float32 => dyn_prim!(f32, f32, array, options, type_),
                 DataType::Float64 => dyn_prim!(f64, f64, array, options, type_),
                 DataType::Utf8 => {


Test patch:
diff --git a/tests/it/array/mod.rs b/tests/it/array/mod.rs
index ef14b993d2..8b400901d7 100644
--- a/tests/it/array/mod.rs
+++ b/tests/it/array/mod.rs
@@ -99,8 +99,8 @@ fn test_with_validity() {
     assert_eq!(arr_ref, &expected);
 }
 
-// check that `PartialEq` can be derived
-#[derive(PartialEq)]
+// check that we ca derive stuff
+#[derive(PartialEq, Clone, Debug)]
 struct A {
     array: Box<dyn Array>,
 }
diff --git a/tests/it/compute/arithmetics/time.rs b/tests/it/compute/arithmetics/time.rs
index cbe53381be..9f31dd60a9 100644
--- a/tests/it/compute/arithmetics/time.rs
+++ b/tests/it/compute/arithmetics/time.rs
@@ -2,6 +2,7 @@ use arrow2::array::*;
 use arrow2::compute::arithmetics::time::*;
 use arrow2::datatypes::{DataType, TimeUnit};
 use arrow2::scalar::*;
+use arrow2::types::months_days_ns;
 
 #[test]
 fn test_adding_timestamp() {
@@ -317,3 +318,68 @@ fn test_date64() {
 
     assert_eq!(result, expected);
 }
+
+#[test]
+fn test_add_interval() {
+    let timestamp =
+        PrimitiveArray::from_slice([1i64]).to(DataType::Timestamp(TimeUnit::Second, None));
+
+    let interval = months_days_ns::new(0, 1, 0);
+
+    let intervals = PrimitiveArray::from_slice([interval]);
+
+    let expected = PrimitiveArray::from_slice([1i64 + 24 * 60 * 60])
+        .to(DataType::Timestamp(TimeUnit::Second, None));
+
+    let result = add_interval(&timestamp, &intervals).unwrap();
+    assert_eq!(result, expected);
+
+    let result = add_interval_scalar(&timestamp, &Some(interval).into()).unwrap();
+    assert_eq!(result, expected);
+}
+
+#[test]
+fn test_add_interval_offset() {
+    let timestamp = PrimitiveArray::from_slice([1i64]).to(DataType::Timestamp(
+        TimeUnit::Second,
+        Some("+01:00".to_string()),
+    ));
+
+    let interval = months_days_ns::new(0, 1, 0);
+
+    let intervals = PrimitiveArray::from_slice([interval]);
+
+    let expected = PrimitiveArray::from_slice([1i64 + 24 * 60 * 60]).to(DataType::Timestamp(
+        TimeUnit::Second,
+        Some("+01:00".to_string()),
+    ));
+
+    let result = add_interval(&timestamp, &intervals).unwrap();
+    assert_eq!(result, expected);
+
+    let result = add_interval_scalar(&timestamp, &Some(interval).into()).unwrap();
+    assert_eq!(result, expected);
+}
+
+#[cfg(feature = "chrono-tz")]
+#[test]
+fn test_add_interval_tz() {
+    let timestamp = PrimitiveArray::from_slice([1i64]).to(DataType::Timestamp(
+        TimeUnit::Second,
+        Some("GMT".to_string()),
+    ));
+
+    let interval = months_days_ns::new(0, 1, 0);
+    let intervals = PrimitiveArray::from_slice([interval]);
+
+    let expected = PrimitiveArray::from_slice([1i64 + 24 * 60 * 60]).to(DataType::Timestamp(
+        TimeUnit::Second,
+        Some("GMT".to_string()),
+    ));
+
+    let result = add_interval(&timestamp, &intervals).unwrap();
+    assert_eq!(result, expected);
+
+    let result = add_interval_scalar(&timestamp, &Some(interval).into()).unwrap();
+    assert_eq!(result, expected);
+}
diff --git a/tests/it/io/parquet/mod.rs b/tests/it/io/parquet/mod.rs
index 955df23e66..3c8b7bafd6 100644
--- a/tests/it/io/parquet/mod.rs
+++ b/tests/it/io/parquet/mod.rs
@@ -962,8 +962,8 @@ fn integration_read(data: &[u8]) -> Result<IntegrationRead> {
 /// logical types.
 #[test]
 fn arrow_type() -> Result<()> {
-    let dt1 = DataType::Duration(TimeUnit::Second);
-    let array = PrimitiveArray::<i64>::from([Some(1), None, Some(2)]).to(dt1.clone());
+    let array1 = PrimitiveArray::<i64>::from([Some(1), None, Some(2)])
+        .to(DataType::Duration(TimeUnit::Second));
     let array2 = Utf8Array::<i64>::from([Some("a"), None, Some("bb")]);
 
     let indices = PrimitiveArray::from_values((0..3u64).map(|x| x % 2));
@@ -978,23 +978,61 @@ fn arrow_type() -> Result<()> {
         vec![b'a', b'b', b'a', b'c'].into(),
         None,
     );
-    let array5 = DictionaryArray::from_data(indices, Box::new(values));
+    let array5 = DictionaryArray::from_data(indices.clone(), Box::new(values));
+
+    let values = PrimitiveArray::from_slice([1i16, 3]);
+    let array6 = DictionaryArray::from_data(indices.clone(), Box::new(values));
+
+    let values = PrimitiveArray::from_slice([1i64, 3]).to(DataType::Timestamp(
+        TimeUnit::Millisecond,
+        Some("UTC".to_string()),
+    ));
+    let array7 = DictionaryArray::from_data(indices.clone(), Box::new(values));
+
+    let values = PrimitiveArray::from_slice([1.0f64, 3.0]);
+    let array8 = DictionaryArray::from_data(indices.clone(), Box::new(values));
+
+    let values = PrimitiveArray::from_slice([1u8, 3]);
+    let array9 = DictionaryArray::from_data(indices.clone(), Box::new(values));
+
+    let values = PrimitiveArray::from_slice([1u16, 3]);
+    let array10 = DictionaryArray::from_data(indices.clone(), Box::new(values));
+
+    let values = PrimitiveArray::from_slice([1u32, 3]);
+    let array11 = DictionaryArray::from_data(indices.clone(), Box::new(values));
+
+    let values = PrimitiveArray::from_slice([1u64, 3]);
+    let array12 = DictionaryArray::from_data(indices, Box::new(values));
 
     let schema = Schema::from(vec![
-        Field::new("a1", dt1, true),
+        Field::new("a1", array1.data_type().clone(), true),
         Field::new("a2", array2.data_type().clone(), true),
         Field::new("a3", array3.data_type().clone(), true),
         Field::new("a4", array4.data_type().clone(), true),
         Field::new("a5", array5.data_type().clone(), true),
-        Field::new("a6", array5.data_type().clone(), false),
+        Field::new("a5a", array5.data_type().clone(), false),
+        Field::new("a6", array6.data_type().clone(), true),
+        Field::new("a7", array7.data_type().clone(), true),
+        Field::new("a8", array8.data_type().clone(), true),
+        Field::new("a9", array9.data_type().clone(), true),
+        Field::new("a10", array10.data_type().clone(), true),
+        Field::new("a11", array11.data_type().clone(), true),
+        Field::new("a12", array12.data_type().clone(), true),
     ]);
     let batch = Chunk::try_new(vec![
-        array.boxed(),
+        array1.boxed(),
         array2.boxed(),
         array3.boxed(),
         array4.boxed(),
         array5.clone().boxed(),
         array5.boxed(),
+        array6.boxed(),
+        array7.boxed(),
+        array8.boxed(),
+        array9.boxed(),
+        array10.boxed(),
+        array11.boxed(),
+        array12.boxed(),
     ])?;
 
     let r = integration_write(&schema, &[batch.clone()])?;