Skip to content

Commit e28fd0d

Browse files
sdf-jklclaude
andauthored
Add DatePart enum 1-indexed variants (#9965)
# Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. --> - Closes #9964. # Rationale for this change Remove extra computations done on the datafusion side for temporal functions. <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> # What changes are included in this PR? - Added two new `DatePart` variant - Closed some tests lists drift <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> # Are these changes tested? - Unit tests <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? If this PR claims a performance improvement, please include evidence such as benchmark results. --> # Are there any user-facing changes? No <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. If there are any breaking changes to public APIs, please call them out. --> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent c4e154f commit e28fd0d

1 file changed

Lines changed: 80 additions & 8 deletions

File tree

arrow-arith/src/temporal.rs

Lines changed: 80 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,10 @@ pub enum DatePart {
6363
DayOfWeekSunday0,
6464
/// Day of the week, in range `0..=6`, where Monday is `0`
6565
DayOfWeekMonday0,
66+
/// Day of the week, in range `1..=7`, where Sunday is `1`
67+
DayOfWeekSunday1,
68+
/// ISO day of the week, in range `1..=7`, where Monday is `1`
69+
DayOfWeekMonday1,
6670
/// Day of year, in range `1..=366`
6771
DayOfYear,
6872
/// Hour of the day, in range `0..=23`
@@ -100,11 +104,6 @@ impl std::fmt::Display for DatePart {
100104
/// - `century`, `decade`, `millennium` — no matching [`DatePart`] variant.
101105
/// - `timezone`, `timezone_hour`, `timezone_minute` — not modeled by
102106
/// [`DatePart`].
103-
/// - `isodow` — PostgreSQL's `isodow` returns `1..=7` (Mon=1), but the
104-
/// closest variant ([`DatePart::DayOfWeekMonday0`]) returns `0..=6`.
105-
/// Accepting it here would silently shift the value range; callers that
106-
/// need the PostgreSQL semantic should map the alias themselves and
107-
/// add `1` to the extracted value.
108107
impl FromStr for DatePart {
109108
type Err = ArrowError;
110109

@@ -118,6 +117,9 @@ impl FromStr for DatePart {
118117
"isoweek" => Self::WeekISO,
119118
"d" | "day" | "days" => Self::Day,
120119
"dow" | "dayofweek" => Self::DayOfWeekSunday0,
120+
"dow1" | "dayofweek1" => Self::DayOfWeekSunday1,
121+
"isodow0" => Self::DayOfWeekMonday0,
122+
"isodow" => Self::DayOfWeekMonday1,
121123
"doy" | "dayofyear" => Self::DayOfYear,
122124
"h" | "hr" | "hrs" | "hour" | "hours" => Self::Hour,
123125
"m" | "min" | "mins" | "minute" | "minutes" => Self::Minute,
@@ -158,6 +160,8 @@ where
158160
DatePart::Day => |d| d.day() as i32,
159161
DatePart::DayOfWeekSunday0 => |d| d.num_days_from_sunday(),
160162
DatePart::DayOfWeekMonday0 => |d| d.num_days_from_monday(),
163+
DatePart::DayOfWeekSunday1 => |d| d.num_days_from_sunday() + 1,
164+
DatePart::DayOfWeekMonday1 => |d| d.num_days_from_monday() + 1,
161165
DatePart::DayOfYear => |d| d.ordinal() as i32,
162166
DatePart::Hour => |d| d.hour() as i32,
163167
DatePart::Minute => |d| d.minute() as i32,
@@ -500,6 +504,8 @@ impl ExtractDatePartExt for PrimitiveArray<IntervalYearMonthType> {
500504
| DatePart::Day
501505
| DatePart::DayOfWeekSunday0
502506
| DatePart::DayOfWeekMonday0
507+
| DatePart::DayOfWeekSunday1
508+
| DatePart::DayOfWeekMonday1
503509
| DatePart::DayOfYear
504510
| DatePart::Hour
505511
| DatePart::Minute
@@ -536,6 +542,8 @@ impl ExtractDatePartExt for PrimitiveArray<IntervalDayTimeType> {
536542
| DatePart::Month
537543
| DatePart::DayOfWeekSunday0
538544
| DatePart::DayOfWeekMonday0
545+
| DatePart::DayOfWeekSunday1
546+
| DatePart::DayOfWeekMonday1
539547
| DatePart::DayOfYear => {
540548
return_compute_error_with!(format!("{part} does not support"), self.data_type())
541549
}
@@ -578,6 +586,8 @@ impl ExtractDatePartExt for PrimitiveArray<IntervalMonthDayNanoType> {
578586
| DatePart::YearISO
579587
| DatePart::DayOfWeekSunday0
580588
| DatePart::DayOfWeekMonday0
589+
| DatePart::DayOfWeekSunday1
590+
| DatePart::DayOfWeekMonday1
581591
| DatePart::DayOfYear => {
582592
return_compute_error_with!(format!("{part} does not support"), self.data_type())
583593
}
@@ -610,6 +620,8 @@ impl ExtractDatePartExt for PrimitiveArray<DurationSecondType> {
610620
| DatePart::Month
611621
| DatePart::DayOfWeekSunday0
612622
| DatePart::DayOfWeekMonday0
623+
| DatePart::DayOfWeekSunday1
624+
| DatePart::DayOfWeekMonday1
613625
| DatePart::DayOfYear => {
614626
return_compute_error_with!(format!("{part} does not support"), self.data_type())
615627
}
@@ -642,6 +654,8 @@ impl ExtractDatePartExt for PrimitiveArray<DurationMillisecondType> {
642654
| DatePart::Month
643655
| DatePart::DayOfWeekSunday0
644656
| DatePart::DayOfWeekMonday0
657+
| DatePart::DayOfWeekSunday1
658+
| DatePart::DayOfWeekMonday1
645659
| DatePart::DayOfYear => {
646660
return_compute_error_with!(format!("{part} does not support"), self.data_type())
647661
}
@@ -674,6 +688,8 @@ impl ExtractDatePartExt for PrimitiveArray<DurationMicrosecondType> {
674688
| DatePart::Month
675689
| DatePart::DayOfWeekSunday0
676690
| DatePart::DayOfWeekMonday0
691+
| DatePart::DayOfWeekSunday1
692+
| DatePart::DayOfWeekMonday1
677693
| DatePart::DayOfYear => {
678694
return_compute_error_with!(format!("{part} does not support"), self.data_type())
679695
}
@@ -706,6 +722,8 @@ impl ExtractDatePartExt for PrimitiveArray<DurationNanosecondType> {
706722
| DatePart::Month
707723
| DatePart::DayOfWeekSunday0
708724
| DatePart::DayOfWeekMonday0
725+
| DatePart::DayOfWeekSunday1
726+
| DatePart::DayOfWeekMonday1
709727
| DatePart::DayOfYear => {
710728
return_compute_error_with!(format!("{part} does not support"), self.data_type())
711729
}
@@ -985,6 +1003,46 @@ mod tests {
9851003
assert_eq!(2, b.value(2));
9861004
}
9871005

1006+
#[test]
1007+
fn test_temporal_array_date64_dayofweek1() {
1008+
//1514764800000 -> 2018-01-01 (Monday)
1009+
//1550636625000 -> 2019-02-20 (Wednesday)
1010+
//1483228800000 -> 2017-01-01 (Sunday)
1011+
let a: PrimitiveArray<Date64Type> = vec![
1012+
Some(1514764800000),
1013+
None,
1014+
Some(1550636625000),
1015+
Some(1483228800000),
1016+
]
1017+
.into();
1018+
1019+
let b = date_part_primitive(&a, DatePart::DayOfWeekSunday1).unwrap();
1020+
assert_eq!(2, b.value(0));
1021+
assert!(!b.is_valid(1));
1022+
assert_eq!(4, b.value(2));
1023+
assert_eq!(1, b.value(3));
1024+
}
1025+
1026+
#[test]
1027+
fn test_temporal_array_date64_isodow() {
1028+
//1514764800000 -> 2018-01-01 (Monday)
1029+
//1550636625000 -> 2019-02-20 (Wednesday)
1030+
//1483228800000 -> 2017-01-01 (Sunday)
1031+
let a: PrimitiveArray<Date64Type> = vec![
1032+
Some(1514764800000),
1033+
None,
1034+
Some(1550636625000),
1035+
Some(1483228800000),
1036+
]
1037+
.into();
1038+
1039+
let b = date_part_primitive(&a, DatePart::DayOfWeekMonday1).unwrap();
1040+
assert_eq!(1, b.value(0));
1041+
assert!(!b.is_valid(1));
1042+
assert_eq!(3, b.value(2));
1043+
assert_eq!(7, b.value(3));
1044+
}
1045+
9881046
#[test]
9891047
fn test_temporal_array_date64_weekday0() {
9901048
//1483228800000 -> 2017-01-01 (Sunday)
@@ -1579,11 +1637,15 @@ mod tests {
15791637
let invalid_parts = [
15801638
DatePart::Quarter,
15811639
DatePart::Year,
1640+
DatePart::YearISO,
15821641
DatePart::Month,
15831642
DatePart::Week,
1643+
DatePart::WeekISO,
15841644
DatePart::Day,
15851645
DatePart::DayOfWeekSunday0,
15861646
DatePart::DayOfWeekMonday0,
1647+
DatePart::DayOfWeekSunday1,
1648+
DatePart::DayOfWeekMonday1,
15871649
DatePart::DayOfYear,
15881650
];
15891651

@@ -1813,8 +1875,12 @@ mod tests {
18131875
fn ensure_returns_error(array: &dyn Array) {
18141876
let invalid_parts = [
18151877
DatePart::Quarter,
1878+
DatePart::YearISO,
1879+
DatePart::WeekISO,
18161880
DatePart::DayOfWeekSunday0,
18171881
DatePart::DayOfWeekMonday0,
1882+
DatePart::DayOfWeekSunday1,
1883+
DatePart::DayOfWeekMonday1,
18181884
DatePart::DayOfYear,
18191885
];
18201886

@@ -1956,10 +2022,14 @@ mod tests {
19562022
fn ensure_returns_error(array: &dyn Array) {
19572023
let invalid_parts = [
19582024
DatePart::Year,
2025+
DatePart::YearISO,
19592026
DatePart::Quarter,
19602027
DatePart::Month,
2028+
DatePart::WeekISO,
19612029
DatePart::DayOfWeekSunday0,
19622030
DatePart::DayOfWeekMonday0,
2031+
DatePart::DayOfWeekSunday1,
2032+
DatePart::DayOfWeekMonday1,
19632033
DatePart::DayOfYear,
19642034
];
19652035

@@ -2157,6 +2227,11 @@ mod tests {
21572227
("day", DatePart::Day),
21582228
("dow", DatePart::DayOfWeekSunday0),
21592229
("DayOfWeek", DatePart::DayOfWeekSunday0),
2230+
("dow1", DatePart::DayOfWeekSunday1),
2231+
("dayofweek1", DatePart::DayOfWeekSunday1),
2232+
("DAYOFWEEK1", DatePart::DayOfWeekSunday1),
2233+
("isodow", DatePart::DayOfWeekMonday1),
2234+
("ISODOW", DatePart::DayOfWeekMonday1),
21602235
("doy", DatePart::DayOfYear),
21612236
("DayOfYear", DatePart::DayOfYear),
21622237
("h", DatePart::Hour),
@@ -2195,8 +2270,6 @@ mod tests {
21952270
#[test]
21962271
fn test_date_part_from_str_unknown() {
21972272
// Names intentionally rejected — see the FromStr doc comment for why.
2198-
// `isodow` is here because mapping it to `DayOfWeekMonday0` would
2199-
// silently shift the value range vs. PostgreSQL's `isodow` (1..=7).
22002273
let unknown = [
22012274
"epoch",
22022275
"century",
@@ -2205,7 +2278,6 @@ mod tests {
22052278
"timezone",
22062279
"timezone_hour",
22072280
"timezone_minute",
2208-
"isodow",
22092281
// Whitespace is not trimmed — pin this so the behavior doesn't
22102282
// change silently. Callers should trim before parsing.
22112283
" year ",

0 commit comments

Comments
 (0)