Skip to content

Commit ad2beef

Browse files
codetyri0ncodetyri0n
committed
minor refactoring and test addition
1 parent 85d7c0d commit ad2beef

File tree

4 files changed

+80
-64
lines changed

4 files changed

+80
-64
lines changed

‎datafusion/functions/src/datetime/common.rs‎

Lines changed: 56 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,26 +17,78 @@
1717

1818
use std::sync::Arc;
1919

20+
use arrow::array::timezone::Tz;
2021
use arrow::array::{
2122
Array, ArrowPrimitiveType, AsArray, GenericStringArray, PrimitiveArray,
2223
StringArrayType, StringViewArray,
2324
};
2425
use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
25-
use arrow::datatypes::DataType;
26+
use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
27+
use arrow::datatypes::{ArrowTimestampType, DataType};
2628
use chrono::format::{parse, Parsed, StrftimeItems};
2729
use chrono::LocalResult::Single;
28-
use chrono::{DateTime, TimeZone, Utc};
30+
use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
31+
use std::ops::Add;
2932

3033
use datafusion_common::cast::as_generic_string_array;
3134
use datafusion_common::{
32-
exec_datafusion_err, exec_err, unwrap_or_internal_err, DataFusionError, Result,
33-
ScalarType, ScalarValue,
35+
exec_datafusion_err, exec_err, internal_datafusion_err, unwrap_or_internal_err,
36+
DataFusionError, Result, ScalarType, ScalarValue,
3437
};
3538
use datafusion_expr::ColumnarValue;
3639

3740
/// Error message if nanosecond conversion request beyond supported interval
3841
const ERR_NANOSECONDS_NOT_SUPPORTED: &str = "The dates that can be represented as nanoseconds have to be between 1677-09-21T00:12:44.0 and 2262-04-11T23:47:16.854775804";
3942

43+
/// Adjusts a timestamp to local time by applying the timezone offset.
44+
pub fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
45+
fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
46+
where
47+
F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
48+
{
49+
match converter(ts) {
50+
MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
51+
"Ambiguous timestamp. Do you mean {:?} or {:?}",
52+
earliest,
53+
latest
54+
),
55+
MappedLocalTime::None => exec_err!(
56+
"The local time does not exist because there is a gap in the local time."
57+
),
58+
MappedLocalTime::Single(date_time) => Ok(date_time),
59+
}
60+
}
61+
62+
let date_time = match T::UNIT {
63+
Nanosecond => Utc.timestamp_nanos(ts),
64+
Microsecond => convert_timestamp(ts, |ts| Utc.timestamp_micros(ts))?,
65+
Millisecond => convert_timestamp(ts, |ts| Utc.timestamp_millis_opt(ts))?,
66+
Second => convert_timestamp(ts, |ts| Utc.timestamp_opt(ts, 0))?,
67+
};
68+
69+
let offset_seconds: i64 = tz
70+
.offset_from_utc_datetime(&date_time.naive_utc())
71+
.fix()
72+
.local_minus_utc() as i64;
73+
74+
let adjusted_date_time = date_time.add(
75+
TimeDelta::try_seconds(offset_seconds)
76+
.ok_or_else(|| internal_datafusion_err!("Offset seconds should be less than i64::MAX / 1_000 or greater than -i64::MAX / 1_000"))?,
77+
);
78+
79+
// convert back to i64
80+
match T::UNIT {
81+
Nanosecond => adjusted_date_time.timestamp_nanos_opt().ok_or_else(|| {
82+
internal_datafusion_err!(
83+
"Failed to convert DateTime to timestamp in nanosecond. This error may occur if the date is out of range. The supported date ranges are between 1677-09-21T00:12:43.145224192 and 2262-04-11T23:47:16.854775807"
84+
)
85+
}),
86+
Microsecond => Ok(adjusted_date_time.timestamp_micros()),
87+
Millisecond => Ok(adjusted_date_time.timestamp_millis()),
88+
Second => Ok(adjusted_date_time.timestamp()),
89+
}
90+
}
91+
4092
/// Calls string_to_timestamp_nanos and converts the error type
4193
pub(crate) fn string_to_timestamp_nanos_shim(s: &str) -> Result<i64> {
4294
string_to_timestamp_nanos(s).map_err(|e| e.into())

‎datafusion/functions/src/datetime/date_part.rs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,7 @@ impl ScalarUDFImpl for DatePartFunc {
207207
};
208208

209209
let part_trim = part_normalization(&part);
210-
let is_epoch = is_epoch(&part);
210+
let is_epoch = is_epoch(part_trim);
211211

212212
// Epoch is timezone-independent - it always returns seconds since 1970-01-01 UTC
213213
let array = if is_epoch {
@@ -333,10 +333,11 @@ fn adjust_timestamp_array<T: ArrowTimestampType>(
333333
}
334334

335335
fn is_epoch(part: &str) -> bool {
336-
let part = part_normalization(part);
337336
matches!(part.to_lowercase().as_str(), "epoch")
338337
}
339338

339+
// Try to remove quote if exist, if the quote is invalid, return original string
340+
// and let the downstream function handle the error.
340341
fn part_normalization(part: &str) -> &str {
341342
part.strip_prefix(|c| c == '\'' || c == '\"')
342343
.and_then(|s| s.strip_suffix(|c| c == '\'' || c == '\"'))
@@ -351,6 +352,7 @@ fn interpret_session_timezone(tz_str: &str) -> Result<Tz> {
351352
}
352353

353354
fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result<ArrayRef> {
355+
// Nanosecond is neither supported in Postgres nor DuckDB, to avoid dealing
354356
// with overflow and precision issue we don't support nanosecond
355357
if unit == Nanosecond {
356358
return not_impl_err!("Date part {unit:?} not supported");

‎datafusion/functions/src/datetime/mod.rs‎

Lines changed: 1 addition & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -19,16 +19,10 @@
1919
2020
use std::sync::Arc;
2121

22-
use arrow::array::timezone::Tz;
23-
use arrow::datatypes::ArrowTimestampType;
24-
use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second};
25-
use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
26-
use datafusion_common::{exec_err, internal_datafusion_err, Result};
27-
use std::ops::Add;
28-
2922
use datafusion_expr::ScalarUDF;
3023

3124
pub mod common;
25+
pub use common::adjust_to_local_time;
3226
pub mod current_date;
3327
pub mod current_time;
3428
pub mod date_bin;
@@ -44,55 +38,6 @@ pub mod to_local_time;
4438
pub mod to_timestamp;
4539
pub mod to_unixtime;
4640

47-
// Adjusts a timestamp to local time by applying the timezone offset.
48-
pub fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
49-
fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
50-
where
51-
F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
52-
{
53-
match converter(ts) {
54-
MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
55-
"Ambiguous timestamp. Do you mean {:?} or {:?}",
56-
earliest,
57-
latest
58-
),
59-
MappedLocalTime::None => exec_err!(
60-
"The local time does not exist because there is a gap in the local time."
61-
),
62-
MappedLocalTime::Single(date_time) => Ok(date_time),
63-
}
64-
}
65-
66-
let date_time = match T::UNIT {
67-
Nanosecond => Utc.timestamp_nanos(ts),
68-
Microsecond => convert_timestamp(ts, |ts| Utc.timestamp_micros(ts))?,
69-
Millisecond => convert_timestamp(ts, |ts| Utc.timestamp_millis_opt(ts))?,
70-
Second => convert_timestamp(ts, |ts| Utc.timestamp_opt(ts, 0))?,
71-
};
72-
73-
let offset_seconds: i64 = tz
74-
.offset_from_utc_datetime(&date_time.naive_utc())
75-
.fix()
76-
.local_minus_utc() as i64;
77-
78-
let adjusted_date_time = date_time.add(
79-
TimeDelta::try_seconds(offset_seconds)
80-
.ok_or_else(|| internal_datafusion_err!("Offset seconds should be less than i64::MAX / 1_000 or greater than -i64::MAX / 1_000"))?,
81-
);
82-
83-
// convert back to i64
84-
match T::UNIT {
85-
Nanosecond => adjusted_date_time.timestamp_nanos_opt().ok_or_else(|| {
86-
internal_datafusion_err!(
87-
"Failed to convert DateTime to timestamp in nanosecond. This error may occur if the date is out of range. The supported date ranges are between 1677-09-21T00:12:43.145224192 and 2262-04-11T23:47:16.854775807"
88-
)
89-
}),
90-
Microsecond => Ok(adjusted_date_time.timestamp_micros()),
91-
Millisecond => Ok(adjusted_date_time.timestamp_millis()),
92-
Second => Ok(adjusted_date_time.timestamp()),
93-
}
94-
}
95-
9641
// create UDFs
9742
make_udf_function!(current_date::CurrentDateFunc, current_date);
9843
make_udf_function!(current_time::CurrentTimeFunc, current_time);

‎datafusion/sqllogictest/test_files/extract_tz.slt‎

Lines changed: 19 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,7 @@ statement ok
127127
SET datafusion.execution.time_zone = 'Asia/Kolkata';
128128

129129
query IIII
130-
SELECT
131-
EXTRACT(HOUR FROM TIMESTAMP '2025-11-22 15:30:45'),
130+
SELECT EXTRACT(HOUR FROM TIMESTAMP '2025-11-22 15:30:45'),
132131
EXTRACT(MINUTE FROM TIMESTAMP '2025-11-22 15:30:45'),
133132
EXTRACT(DOW FROM TIMESTAMP '2025-11-22 00:00:00'),
134133
EXTRACT(SECOND FROM TIMESTAMP '2024-01-01 03:05:59');
@@ -156,3 +155,21 @@ query I
156155
SELECT EXTRACT(HOUR FROM TIMESTAMP '2025-01-15 10:00:00');
157156
----
158157
5
158+
159+
statement ok
160+
SET datafusion.execution.time_zone = '-03:30';
161+
162+
query II
163+
SELECT EXTRACT(MINUTE FROM TIMESTAMP '2023-10-30 10:45:30'),
164+
EXTRACT(SECOND FROM TIMESTAMP '2023-10-30 10:45:30');
165+
----
166+
15 30
167+
168+
statement ok
169+
SET datafusion.execution.time_zone = 'America/St_Johns';
170+
171+
query II
172+
SELECT EXTRACT(MINUTE FROM TIMESTAMP '2023-10-30 10:45:30'),
173+
EXTRACT(SECOND FROM TIMESTAMP '2023-10-30 10:45:30');
174+
----
175+
15 30

0 commit comments

Comments
 (0)