Fixed panic in writing dictionaries to parquet

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,

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 => {

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>,
 }

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);
+}

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()])?;