Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Commit

Permalink
Added support for MONTH_DAY_NANO intervals.
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao committed Aug 13, 2021
1 parent 86edfd6 commit be7d991
Show file tree
Hide file tree
Showing 16 changed files with 353 additions and 49 deletions.
7 changes: 6 additions & 1 deletion src/array/equal/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use std::unimplemented;

use crate::{
datatypes::{DataType, IntervalUnit},
types::{days_ms, NativeType},
types::{days_ms, months_days_ns, NativeType},
};

use super::{
Expand Down Expand Up @@ -224,6 +224,11 @@ pub fn equal(lhs: &dyn Array, rhs: &dyn Array) -> bool {
let rhs = rhs.as_any().downcast_ref().unwrap();
primitive::equal::<days_ms>(lhs, rhs)
}
DataType::Interval(IntervalUnit::MonthDayNano) => {
let lhs = lhs.as_any().downcast_ref().unwrap();
let rhs = rhs.as_any().downcast_ref().unwrap();
primitive::equal::<months_days_ns>(lhs, rhs)
}
DataType::Float16 => unreachable!(),
DataType::Float32 => {
let lhs = lhs.as_any().downcast_ref().unwrap();
Expand Down
3 changes: 3 additions & 0 deletions src/array/ffi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,9 @@ pub fn buffers_children_dictionary(array: &dyn Array) -> BuffersChildren {
ffi_dyn!(array, PrimitiveArray<i32>)
}
DataType::Interval(IntervalUnit::DayTime) => ffi_dyn!(array, PrimitiveArray<days_ms>),
DataType::Interval(IntervalUnit::MonthDayNano) => {
ffi_dyn!(array, PrimitiveArray<months_days_ns>)
}
DataType::Int64
| DataType::Date64
| DataType::Time64(_)
Expand Down
3 changes: 3 additions & 0 deletions src/array/growable/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,9 @@ pub fn make_growable<'a>(
DataType::Interval(IntervalUnit::DayTime) => {
dyn_growable!(days_ms, arrays, use_validity, capacity)
}
DataType::Interval(IntervalUnit::MonthDayNano) => {
dyn_growable!(months_days_ns, arrays, use_validity, capacity)
}
DataType::Decimal(_, _) => dyn_growable!(i128, arrays, use_validity, capacity),
DataType::UInt8 => dyn_growable!(u8, arrays, use_validity, capacity),
DataType::UInt16 => dyn_growable!(u16, arrays, use_validity, capacity),
Expand Down
14 changes: 13 additions & 1 deletion src/array/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,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, IntervalUnit},
Expand Down Expand Up @@ -163,6 +163,9 @@ impl Display for dyn Array {
DataType::Interval(IntervalUnit::DayTime) => {
fmt_dyn!(self, PrimitiveArray<days_ms>, f)
}
DataType::Interval(IntervalUnit::MonthDayNano) => {
fmt_dyn!(self, PrimitiveArray<months_days_ns>, f)
}
DataType::Int64
| DataType::Date64
| DataType::Time64(_)
Expand Down Expand Up @@ -217,6 +220,9 @@ pub fn new_empty_array(data_type: DataType) -> Box<dyn Array> {
DataType::Interval(IntervalUnit::DayTime) => {
Box::new(PrimitiveArray::<days_ms>::new_empty(data_type))
}
DataType::Interval(IntervalUnit::MonthDayNano) => {
Box::new(PrimitiveArray::<months_days_ns>::new_empty(data_type))
}
DataType::Int64
| DataType::Date64
| DataType::Time64(_)
Expand Down Expand Up @@ -271,6 +277,9 @@ pub fn new_null_array(data_type: DataType, length: usize) -> Box<dyn Array> {
DataType::Interval(IntervalUnit::DayTime) => {
Box::new(PrimitiveArray::<days_ms>::new_null(data_type, length))
}
DataType::Interval(IntervalUnit::MonthDayNano) => Box::new(
PrimitiveArray::<months_days_ns>::new_null(data_type, length),
),
DataType::Int64
| DataType::Date64
| DataType::Time64(_)
Expand Down Expand Up @@ -332,6 +341,9 @@ pub fn clone(array: &dyn Array) -> Box<dyn Array> {
clone_dyn!(array, PrimitiveArray<i32>)
}
DataType::Interval(IntervalUnit::DayTime) => clone_dyn!(array, PrimitiveArray<days_ms>),
DataType::Interval(IntervalUnit::MonthDayNano) => {
clone_dyn!(array, PrimitiveArray<months_days_ns>)
}
DataType::Int64
| DataType::Date64
| DataType::Time64(_)
Expand Down
18 changes: 17 additions & 1 deletion src/array/primitive/display.rs
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
use crate::{datatypes::*, temporal_conversions, types::days_ms};
use crate::{
datatypes::*,
temporal_conversions,
types::{days_ms, months_days_ns},
};

use super::super::{display_fmt, Array};
use super::PrimitiveArray;
Expand Down Expand Up @@ -159,6 +163,18 @@ impl std::fmt::Display for PrimitiveArray<days_ms> {
}
}

impl std::fmt::Display for PrimitiveArray<months_days_ns> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let new_lines = false;
let head = &format!("{}", self.data_type());
let iter = self.iter().map(|x| {
x.copied()
.map(|x| format!("{}m{}d{}ns", x.months(), x.days(), x.ns()))
});
display_fmt(iter, head, f, new_lines)
}
}

macro_rules! display {
($ty:ty) => {
impl std::fmt::Display for PrimitiveArray<$ty> {
Expand Down
6 changes: 5 additions & 1 deletion src/array/primitive/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -167,6 +167,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`
Expand All @@ -192,6 +194,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`
Expand Down
3 changes: 2 additions & 1 deletion src/compute/aggregate/memory.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use crate::array::*;
use crate::bitmap::Bitmap;
use crate::datatypes::{DataType, IntervalUnit};
use crate::types::days_ms;
use crate::types::{days_ms, months_days_ns};

fn validity_size(validity: &Option<Bitmap>) -> usize {
validity.as_ref().map(|b| b.as_slice().len()).unwrap_or(0)
Expand Down Expand Up @@ -72,6 +72,7 @@ pub fn estimated_bytes_size(array: &dyn Array) -> usize {
Float64 => dyn_primitive!(array, f64),
Decimal(_, _) => dyn_primitive!(array, i128),
Interval(IntervalUnit::DayTime) => dyn_primitive!(array, days_ms),
Interval(IntervalUnit::MonthDayNano) => dyn_primitive!(array, months_days_ns),
Binary => dyn_binary!(array, BinaryArray<i32>, i32),
FixedSizeBinary(_) => {
let array = array
Expand Down
7 changes: 7 additions & 0 deletions src/datatypes/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,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 {
Expand Down
2 changes: 2 additions & 0 deletions src/io/ipc/convert.rs
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,7 @@ pub(crate) fn get_data_type(field: ipc::Field, may_be_dictionary: bool) -> DataT
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),
}
}
Expand Down Expand Up @@ -513,6 +514,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 {
Expand Down
Loading

0 comments on commit be7d991

Please sign in to comment.