diff --git a/src/io/parquet/read/deserialize/simple.rs b/src/io/parquet/read/deserialize/simple.rs index aaab8aee9f5..20dba380653 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::::new( + iter, + data_type, + chunk_size, + |x: i64| x as u64, + )), Int8 => dyn_iter(primitive::DictIter::::new( iter, data_type, diff --git a/src/io/parquet/write/dictionary.rs b/src/io/parquet/write/dictionary.rs index 65f13a49042..05c56207bfe 100644 --- a/src/io/parquet/write/dictionary.rs +++ b/src/io/parquet/write/dictionary.rs @@ -149,7 +149,7 @@ pub fn array_to_pages( 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 => { diff --git a/tests/it/array/mod.rs b/tests/it/array/mod.rs index ef14b993d25..8b400901d70 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, } diff --git a/tests/it/compute/arithmetics/time.rs b/tests/it/compute/arithmetics/time.rs index cbe53381bef..9f31dd60a97 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(×tamp, &intervals).unwrap(); + assert_eq!(result, expected); + + let result = add_interval_scalar(×tamp, &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(×tamp, &intervals).unwrap(); + assert_eq!(result, expected); + + let result = add_interval_scalar(×tamp, &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(×tamp, &intervals).unwrap(); + assert_eq!(result, expected); + + let result = add_interval_scalar(×tamp, &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 955df23e663..3c8b7bafd6d 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 { /// logical types. #[test] fn arrow_type() -> Result<()> { - let dt1 = DataType::Duration(TimeUnit::Second); - let array = PrimitiveArray::::from([Some(1), None, Some(2)]).to(dt1.clone()); + let array1 = PrimitiveArray::::from([Some(1), None, Some(2)]) + .to(DataType::Duration(TimeUnit::Second)); let array2 = Utf8Array::::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()])?;