Skip to content

Commit 6d6c939

Browse files
kosiewkosiew
authored
Add Week ISO, Year ISO computation (#7163)
* feat: Add WeekISO alias to DatePart enum and update date_part function * feat: Add YearISO variant to DatePart enum and implement its functionality in date_part function * fix: improve ISO week , year calculations in date_part functions * Revert "fix: improve ISO week , year calculations in date_part functions" This reverts commit 1a85be4. * test: add ISO week and year calculations for date64 and timestamp arrays * fix: add debug logging for year and ISO year calculations in date_part functions * test: add constants for readability * move WeekISO, YearISO to unsupported for Interval, Duration.. * fix: update random number generation to use thread_rng and gen_range * fix: clarify ISO year and week documentation in DatePart enum
1 parent f6ac87e commit 6d6c939

File tree

1 file changed

+155
-5
lines changed

1 file changed

+155
-5
lines changed

‎arrow-arith/src/temporal.rs

Lines changed: 155 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -47,10 +47,14 @@ pub enum DatePart {
4747
Quarter,
4848
/// Calendar year
4949
Year,
50+
/// ISO year, computed as per ISO 8601
51+
YearISO,
5052
/// Month in the year, in range `1..=12`
5153
Month,
52-
/// ISO week of the year, in range `1..=53`
54+
/// week of the year, in range `1..=53`, computed as per ISO 8601
5355
Week,
56+
/// ISO week of the year, in range `1..=53`
57+
WeekISO,
5458
/// Day of the month, in range `1..=31`
5559
Day,
5660
/// Day of the week, in range `0..=6`, where Sunday is `0`
@@ -91,8 +95,9 @@ where
9195
match part {
9296
DatePart::Quarter => |d| d.quarter() as i32,
9397
DatePart::Year => |d| d.year(),
98+
DatePart::YearISO => |d| d.iso_week().year(),
9499
DatePart::Month => |d| d.month() as i32,
95-
DatePart::Week => |d| d.iso_week().week() as i32,
100+
DatePart::Week | DatePart::WeekISO => |d| d.iso_week().week() as i32,
96101
DatePart::Day => |d| d.day() as i32,
97102
DatePart::DayOfWeekSunday0 => |d| d.num_days_from_sunday(),
98103
DatePart::DayOfWeekMonday0 => |d| d.num_days_from_monday(),
@@ -102,7 +107,7 @@ where
102107
DatePart::Second => |d| d.second() as i32,
103108
DatePart::Millisecond => |d| (d.nanosecond() / 1_000_000) as i32,
104109
DatePart::Microsecond => |d| (d.nanosecond() / 1_000) as i32,
105-
DatePart::Nanosecond => |d| (d.nanosecond()) as i32,
110+
DatePart::Nanosecond => |d| d.nanosecond() as i32,
106111
}
107112
}
108113

@@ -130,9 +135,14 @@ where
130135
/// let input: TimestampMicrosecondArray =
131136
/// vec![Some(1612025847000000), None, Some(1722015847000000)].into();
132137
///
133-
/// let actual = date_part(&input, DatePart::Week).unwrap();
138+
/// let week = date_part(&input, DatePart::Week).unwrap();
139+
/// let week_iso = date_part(&input, DatePart::WeekISO).unwrap();
134140
/// let expected: Int32Array = vec![Some(4), None, Some(30)].into();
135-
/// assert_eq!(actual.as_ref(), &expected);
141+
/// assert_eq!(week.as_ref(), &expected);
142+
/// assert_eq!(week_iso.as_ref(), &expected);
143+
/// let year_iso = date_part(&input, DatePart::YearISO).unwrap();
144+
/// let expected: Int32Array = vec![Some(2021), None, Some(2024)].into();
145+
/// assert_eq!(year_iso.as_ref(), &expected);
136146
/// ```
137147
pub fn date_part(array: &dyn Array, part: DatePart) -> Result<ArrayRef, ArrowError> {
138148
downcast_temporal_array!(
@@ -430,6 +440,8 @@ impl ExtractDatePartExt for PrimitiveArray<IntervalYearMonthType> {
430440

431441
DatePart::Quarter
432442
| DatePart::Week
443+
| DatePart::WeekISO
444+
| DatePart::YearISO
433445
| DatePart::Day
434446
| DatePart::DayOfWeekSunday0
435447
| DatePart::DayOfWeekMonday0
@@ -460,6 +472,8 @@ impl ExtractDatePartExt for PrimitiveArray<IntervalDayTimeType> {
460472

461473
DatePart::Quarter
462474
| DatePart::Year
475+
| DatePart::YearISO
476+
| DatePart::WeekISO
463477
| DatePart::Month
464478
| DatePart::DayOfWeekSunday0
465479
| DatePart::DayOfWeekMonday0
@@ -495,6 +509,8 @@ impl ExtractDatePartExt for PrimitiveArray<IntervalMonthDayNanoType> {
495509
DatePart::Nanosecond => Ok(self.unary_opt(|d| d.nanoseconds.try_into().ok())),
496510

497511
DatePart::Quarter
512+
| DatePart::WeekISO
513+
| DatePart::YearISO
498514
| DatePart::DayOfWeekSunday0
499515
| DatePart::DayOfWeekMonday0
500516
| DatePart::DayOfYear => {
@@ -523,6 +539,8 @@ impl ExtractDatePartExt for PrimitiveArray<DurationSecondType> {
523539
),
524540

525541
DatePart::Year
542+
| DatePart::YearISO
543+
| DatePart::WeekISO
526544
| DatePart::Quarter
527545
| DatePart::Month
528546
| DatePart::DayOfWeekSunday0
@@ -553,6 +571,8 @@ impl ExtractDatePartExt for PrimitiveArray<DurationMillisecondType> {
553571
}
554572

555573
DatePart::Year
574+
| DatePart::YearISO
575+
| DatePart::WeekISO
556576
| DatePart::Quarter
557577
| DatePart::Month
558578
| DatePart::DayOfWeekSunday0
@@ -583,6 +603,8 @@ impl ExtractDatePartExt for PrimitiveArray<DurationMicrosecondType> {
583603
}
584604

585605
DatePart::Year
606+
| DatePart::YearISO
607+
| DatePart::WeekISO
586608
| DatePart::Quarter
587609
| DatePart::Month
588610
| DatePart::DayOfWeekSunday0
@@ -613,6 +635,8 @@ impl ExtractDatePartExt for PrimitiveArray<DurationNanosecondType> {
613635
DatePart::Nanosecond => Ok(self.unary_opt(|d| d.try_into().ok())),
614636

615637
DatePart::Year
638+
| DatePart::YearISO
639+
| DatePart::WeekISO
616640
| DatePart::Quarter
617641
| DatePart::Month
618642
| DatePart::DayOfWeekSunday0
@@ -2072,4 +2096,130 @@ mod tests {
20722096
ensure_returns_error(&DurationMicrosecondArray::from(vec![0]));
20732097
ensure_returns_error(&DurationNanosecondArray::from(vec![0]));
20742098
}
2099+
2100+
const TIMESTAMP_SECOND_1970_01_01: i64 = 0;
2101+
const TIMESTAMP_SECOND_2018_01_01: i64 = 1_514_764_800;
2102+
const TIMESTAMP_SECOND_2019_02_20: i64 = 1_550_636_625;
2103+
const SECONDS_IN_DAY: i64 = 24 * 60 * 60;
2104+
// In 2018 the ISO year and calendar year start on the same date— 2018-01-01 or 2018-W01-1
2105+
#[test]
2106+
fn test_temporal_array_date64_week_iso() {
2107+
let a: PrimitiveArray<Date64Type> = vec![
2108+
Some(TIMESTAMP_SECOND_2018_01_01 * 1000),
2109+
Some(TIMESTAMP_SECOND_2019_02_20 * 1000),
2110+
]
2111+
.into();
2112+
2113+
let b = date_part(&a, DatePart::WeekISO).unwrap();
2114+
let actual = b.as_primitive::<Int32Type>();
2115+
assert_eq!(1, actual.value(0));
2116+
assert_eq!(8, actual.value(1));
2117+
}
2118+
2119+
#[test]
2120+
fn test_temporal_array_date64_year_iso() {
2121+
let a: PrimitiveArray<Date64Type> = vec![
2122+
Some(TIMESTAMP_SECOND_2018_01_01 * 1000),
2123+
Some(TIMESTAMP_SECOND_2019_02_20 * 1000),
2124+
]
2125+
.into();
2126+
2127+
let b = date_part(&a, DatePart::YearISO).unwrap();
2128+
let actual = b.as_primitive::<Int32Type>();
2129+
assert_eq!(2018, actual.value(0));
2130+
assert_eq!(2019, actual.value(1));
2131+
}
2132+
2133+
#[test]
2134+
fn test_temporal_array_timestamp_week_iso() {
2135+
let a = TimestampSecondArray::from(vec![
2136+
TIMESTAMP_SECOND_1970_01_01, // 0 and is Thursday
2137+
SECONDS_IN_DAY * 4, // Monday of week 2
2138+
SECONDS_IN_DAY * 4 - 1, // Sunday of week 1
2139+
]);
2140+
let b = date_part(&a, DatePart::WeekISO).unwrap();
2141+
let actual = b.as_primitive::<Int32Type>();
2142+
assert_eq!(1, actual.value(0));
2143+
assert_eq!(2, actual.value(1));
2144+
assert_eq!(1, actual.value(2));
2145+
}
2146+
2147+
#[test]
2148+
fn test_temporal_array_timestamp_year_iso() {
2149+
let a = TimestampSecondArray::from(vec![
2150+
TIMESTAMP_SECOND_1970_01_01,
2151+
SECONDS_IN_DAY * 4,
2152+
SECONDS_IN_DAY * 4 - 1,
2153+
]);
2154+
let b = date_part(&a, DatePart::YearISO).unwrap();
2155+
let actual = b.as_primitive::<Int32Type>();
2156+
assert_eq!(1970, actual.value(0));
2157+
assert_eq!(1970, actual.value(1));
2158+
assert_eq!(1970, actual.value(2));
2159+
}
2160+
2161+
const TIMESTAMP_SECOND_2015_12_28: i64 = 1_451_260_800;
2162+
const TIMESTAMP_SECOND_2016_01_03: i64 = 1_451_779_200;
2163+
// January 1st 2016 is a Friday, so 2015 week 53 runs from
2164+
// 2015-12-28 to 2016-01-03 inclusive, and
2165+
// 2016 week 1 runs from 2016-01-04 to 2016-01-10 inclusive.
2166+
#[test]
2167+
fn test_temporal_array_date64_week_iso_edge_cases() {
2168+
let a: PrimitiveArray<Date64Type> = vec![
2169+
Some(TIMESTAMP_SECOND_2015_12_28 * 1000),
2170+
Some(TIMESTAMP_SECOND_2016_01_03 * 1000),
2171+
Some((TIMESTAMP_SECOND_2016_01_03 + SECONDS_IN_DAY) * 1000),
2172+
]
2173+
.into();
2174+
2175+
let b = date_part(&a, DatePart::WeekISO).unwrap();
2176+
let actual = b.as_primitive::<Int32Type>();
2177+
assert_eq!(53, actual.value(0));
2178+
assert_eq!(53, actual.value(1));
2179+
assert_eq!(1, actual.value(2));
2180+
}
2181+
2182+
#[test]
2183+
fn test_temporal_array_date64_year_iso_edge_cases() {
2184+
let a: PrimitiveArray<Date64Type> = vec![
2185+
Some(TIMESTAMP_SECOND_2015_12_28 * 1000),
2186+
Some(TIMESTAMP_SECOND_2016_01_03 * 1000),
2187+
Some((TIMESTAMP_SECOND_2016_01_03 + SECONDS_IN_DAY) * 1000),
2188+
]
2189+
.into();
2190+
2191+
let b = date_part(&a, DatePart::YearISO).unwrap();
2192+
let actual = b.as_primitive::<Int32Type>();
2193+
assert_eq!(2015, actual.value(0));
2194+
assert_eq!(2015, actual.value(1));
2195+
assert_eq!(2016, actual.value(2));
2196+
}
2197+
2198+
#[test]
2199+
fn test_temporal_array_timestamp_week_iso_edge_cases() {
2200+
let a = TimestampSecondArray::from(vec![
2201+
TIMESTAMP_SECOND_2015_12_28,
2202+
TIMESTAMP_SECOND_2016_01_03,
2203+
TIMESTAMP_SECOND_2016_01_03 + SECONDS_IN_DAY,
2204+
]);
2205+
let b = date_part(&a, DatePart::WeekISO).unwrap();
2206+
let actual = b.as_primitive::<Int32Type>();
2207+
assert_eq!(53, actual.value(0));
2208+
assert_eq!(53, actual.value(1));
2209+
assert_eq!(1, actual.value(2));
2210+
}
2211+
2212+
#[test]
2213+
fn test_temporal_array_timestamp_year_iso_edge_cases() {
2214+
let a = TimestampSecondArray::from(vec![
2215+
TIMESTAMP_SECOND_2015_12_28,
2216+
TIMESTAMP_SECOND_2016_01_03,
2217+
TIMESTAMP_SECOND_2016_01_03 + SECONDS_IN_DAY,
2218+
]);
2219+
let b = date_part(&a, DatePart::YearISO).unwrap();
2220+
let actual = b.as_primitive::<Int32Type>();
2221+
assert_eq!(2015, actual.value(0));
2222+
assert_eq!(2015, actual.value(1));
2223+
assert_eq!(2016, actual.value(2));
2224+
}
20752225
}

0 commit comments

Comments
 (0)