Added support for MONTH_DAY_NANO interval type (#268)

* Added support for MONTH_DAY_NANO intervals. * Added integration tests. * Added some tests.
jorgecarleitao · Aug 31, 2021 · 3d7cc31 · 3d7cc31
1 parent 2f1d8fe
commit 3d7cc31
Show file tree

Hide file tree

Showing 22 changed files with 341 additions and 62 deletions.
diff --git a/integration-testing/unskip.patch b/integration-testing/unskip.patch
@@ -1,5 +1,5 @@
 diff --git a/dev/archery/archery/integration/datagen.py b/dev/archery/archery/integration/datagen.py
-index d0c4b3d6c..bc4b83e68 100644
+index d0c4b3d6c..936351c80 100644
 --- a/dev/archery/archery/integration/datagen.py
 +++ b/dev/archery/archery/integration/datagen.py
 @@ -1568,8 +1568,7 @@ def get_generated_json_files(tempdir=None):
@@ -12,7 +12,7 @@ index d0c4b3d6c..bc4b83e68 100644
 
          generate_decimal256_case()
          .skip_category('Go')  # TODO(ARROW-7948): Decimal + Go
-@@ -1579,8 +1578,7 @@ def get_generated_json_files(tempdir=None):
+@@ -1579,13 +1578,11 @@ def get_generated_json_files(tempdir=None):
          generate_datetime_case(),
 
          generate_interval_case()
@@ -22,7 +22,13 @@ index d0c4b3d6c..bc4b83e68 100644
 
          generate_month_day_nano_interval_case()
          .skip_category('Go')
-@@ -1603,13 +1601,11 @@ def get_generated_json_files(tempdir=None):
+-        .skip_category('JS')
+-        .skip_category('Rust'),
++        .skip_category('JS'),
+
+
+         generate_map_case()
+@@ -1603,13 +1600,11 @@ def get_generated_json_files(tempdir=None):
 
          generate_nested_large_offsets_case()
          .skip_category('Go')

diff --git a/src/array/display.rs b/src/array/display.rs
@@ -128,6 +128,15 @@ pub fn get_value_display<'a>(array: &'a dyn Array) -> Box<dyn Fn(usize) -> Strin
                 x.milliseconds()
             ))
         }
+
+        Interval(IntervalUnit::MonthDayNano) => {
+            dyn_primitive!(array, months_days_ns, |x: months_days_ns| format!(
+                "{}m{}d{}ns",
+                x.months(),
+                x.days(),
+                x.ns()
+            ))
+        }
         Duration(TimeUnit::Second) => dyn_primitive!(array, i64, |x| format!("{}s", x)),
         Duration(TimeUnit::Millisecond) => dyn_primitive!(array, i64, |x| format!("{}ms", x)),
         Duration(TimeUnit::Microsecond) => dyn_primitive!(array, i64, |x| format!("{}us", x)),

diff --git a/src/array/equal/mod.rs b/src/array/equal/mod.rs
@@ -1,4 +1,4 @@
-use crate::types::{days_ms, NativeType};
+use crate::types::NativeType;
 
 use super::*;
 

diff --git a/src/array/mod.rs b/src/array/mod.rs
@@ -17,7 +17,7 @@ use std::any::Any;
 use std::fmt::Display;
 
 use crate::error::Result;
-use crate::types::days_ms;
+use crate::types::{days_ms, months_days_ns};
 use crate::{
     bitmap::{Bitmap, MutableBitmap},
     datatypes::DataType,
@@ -186,13 +186,15 @@ macro_rules! with_match_primitive_type {(
 ) => ({
     macro_rules! __with_ty__ {( $_ $T:ident ) => ( $($body)* )}
     use crate::datatypes::PrimitiveType::*;
+    use crate::types::{days_ms, months_days_ns};
     match $key_type {
         Int8 => __with_ty__! { i8 },
         Int16 => __with_ty__! { i16 },
         Int32 => __with_ty__! { i32 },
         Int64 => __with_ty__! { i64 },
         Int128 => __with_ty__! { i128 },
         DaysMs => __with_ty__! { days_ms },
+        MonthDayNano => __with_ty__! { months_days_ns },
         UInt8 => __with_ty__! { u8 },
         UInt16 => __with_ty__! { u16 },
         UInt32 => __with_ty__! { u32 },

diff --git a/src/array/primitive/mod.rs b/src/array/primitive/mod.rs
@@ -3,7 +3,7 @@ use crate::{
     buffer::Buffer,
     datatypes::*,
     error::ArrowError,
-    types::{days_ms, NativeType},
+    types::{days_ms, months_days_ns, NativeType},
 };
 
 use super::Array;
@@ -176,6 +176,8 @@ pub type Int64Array = PrimitiveArray<i64>;
 pub type Int128Array = PrimitiveArray<i128>;
 /// A type definition [`PrimitiveArray`] for [`days_ms`]
 pub type DaysMsArray = PrimitiveArray<days_ms>;
+/// A type definition [`PrimitiveArray`] for [`months_days_ns`]
+pub type MonthsDaysNsArray = PrimitiveArray<months_days_ns>;
 /// A type definition [`PrimitiveArray`] for `f32`
 pub type Float32Array = PrimitiveArray<f32>;
 /// A type definition [`PrimitiveArray`] for `f64`
@@ -201,6 +203,8 @@ pub type Int64Vec = MutablePrimitiveArray<i64>;
 pub type Int128Vec = MutablePrimitiveArray<i128>;
 /// A type definition [`MutablePrimitiveArray`] for [`days_ms`]
 pub type DaysMsVec = MutablePrimitiveArray<days_ms>;
+/// A type definition [`MutablePrimitiveArray`] for [`months_days_ns`]
+pub type MonthsDaysNsVec = MutablePrimitiveArray<months_days_ns>;
 /// A type definition [`MutablePrimitiveArray`] for `f32`
 pub type Float32Vec = MutablePrimitiveArray<f32>;
 /// A type definition [`MutablePrimitiveArray`] for `f64`

diff --git a/src/compute/aggregate/memory.rs b/src/compute/aggregate/memory.rs
@@ -1,7 +1,6 @@
 use crate::array::*;
 use crate::bitmap::Bitmap;
 use crate::datatypes::PhysicalType;
-use crate::types::days_ms;
 
 fn validity_size(validity: &Option<Bitmap>) -> usize {
     validity.as_ref().map(|b| b.as_slice().0.len()).unwrap_or(0)

diff --git a/src/compute/filter.rs b/src/compute/filter.rs
@@ -15,14 +15,10 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::array::growable::make_growable;
-use crate::array::growable::Growable;
+use crate::array::growable::{make_growable, Growable};
+use crate::bitmap::{utils::SlicesIterator, Bitmap, MutableBitmap};
 use crate::record_batch::RecordBatch;
-use crate::{array::*, bitmap::Bitmap, types::NativeType};
-use crate::{
-    bitmap::{utils::SlicesIterator, MutableBitmap},
-    types::days_ms,
-};
+use crate::{array::*, types::NativeType};
 use crate::{buffer::MutableBuffer, error::Result};
 
 /// Function that can filter arbitrary arrays

diff --git a/src/compute/take/mod.rs b/src/compute/take/mod.rs
@@ -21,7 +21,7 @@ use crate::{
     array::{new_empty_array, Array, NullArray, PrimitiveArray},
     datatypes::DataType,
     error::Result,
-    types::{days_ms, Index},
+    types::Index,
 };
 
 mod binary;

diff --git a/src/datatypes/mod.rs b/src/datatypes/mod.rs
@@ -147,6 +147,13 @@ pub enum IntervalUnit {
     /// Indicates the number of elapsed days and milliseconds,
     /// stored as 2 contiguous 32-bit integers (8-bytes in total).
     DayTime,
+    /// The values are stored contiguously in 16 byte blocks. Months and
+    /// days are encoded as 32 bit integers and nanoseconds is encoded as a
+    /// 64 bit integer. All integers are signed. Each field is independent
+    /// (e.g. there is no constraint that nanoseconds have the same sign
+    /// as days or that the quantitiy of nanoseconds represents less
+    /// then a day's worth of time).
+    MonthDayNano,
 }
 
 impl DataType {
@@ -197,6 +204,9 @@ impl DataType {
             Float32 => PhysicalType::Primitive(PrimitiveType::Float32),
             Float64 => PhysicalType::Primitive(PrimitiveType::Float64),
             Interval(IntervalUnit::DayTime) => PhysicalType::Primitive(PrimitiveType::DaysMs),
+            Interval(IntervalUnit::MonthDayNano) => {
+                PhysicalType::Primitive(PrimitiveType::MonthDayNano)
+            }
             Binary => PhysicalType::Binary,
             FixedSizeBinary(_) => PhysicalType::FixedSizeBinary,
             LargeBinary => PhysicalType::LargeBinary,

diff --git a/src/datatypes/physical_type.rs b/src/datatypes/physical_type.rs
@@ -45,6 +45,8 @@ pub enum PrimitiveType {
     Float64,
     /// Two i32 representing days and ms
     DaysMs,
+    /// months_days_ns(i32, i32, i64)
+    MonthDayNano,
 }
 
 /// The set of physical types: unique in-memory representations of an Arrow array.

diff --git a/src/ffi/array.rs b/src/ffi/array.rs
@@ -20,7 +20,6 @@
 use super::ffi::ArrowArrayRef;
 use crate::array::{BooleanArray, FromFfi};
 use crate::error::{ArrowError, Result};
-use crate::types::days_ms;
 use crate::{array::*, datatypes::PhysicalType};
 
 /// Reads a valid `ffi` interface into a `Box<dyn Array>`

diff --git a/src/ffi/schema.rs b/src/ffi/schema.rs
@@ -323,6 +323,9 @@ fn to_format(data_type: &DataType) -> String {
         DataType::Duration(TimeUnit::Nanosecond) => "tDn".to_string(),
         DataType::Interval(IntervalUnit::YearMonth) => "tiM".to_string(),
         DataType::Interval(IntervalUnit::DayTime) => "tiD".to_string(),
+        DataType::Interval(IntervalUnit::MonthDayNano) => {
+            todo!("Spec for FFI for MonthDayNano still not defined.")
+        }
         DataType::Timestamp(unit, tz) => {
             let unit = match unit {
                 TimeUnit::Second => "s".to_string(),

diff --git a/src/io/ipc/convert.rs b/src/io/ipc/convert.rs
@@ -255,6 +255,7 @@ fn get_data_type(field: ipc::Field, extension: Extension, may_be_dictionary: boo
             match interval.unit() {
                 ipc::IntervalUnit::YEAR_MONTH => DataType::Interval(IntervalUnit::YearMonth),
                 ipc::IntervalUnit::DAY_TIME => DataType::Interval(IntervalUnit::DayTime),
+                ipc::IntervalUnit::MONTH_DAY_NANO => DataType::Interval(IntervalUnit::MonthDayNano),
                 z => panic!("Interval type with unit of {:?} unsupported", z),
             }
         }
@@ -604,6 +605,7 @@ pub(crate) fn get_fb_field_type<'a>(
             let interval_unit = match unit {
                 IntervalUnit::YearMonth => ipc::IntervalUnit::YEAR_MONTH,
                 IntervalUnit::DayTime => ipc::IntervalUnit::DAY_TIME,
+                IntervalUnit::MonthDayNano => ipc::IntervalUnit::MONTH_DAY_NANO,
             };
             builder.add_unit(interval_unit);
             FbFieldType {