You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "appletreeisyellow (via GitHub)" <gi...@apache.org> on 2024/07/09 02:56:10 UTC
[PR] feat: add UDF to_local_time() [datafusion]
appletreeisyellow opened a new pull request, #11347:
URL: https://github.com/apache/datafusion/pull/11347
## Which issue does this PR close?
Help with https://github.com/apache/datafusion/issues/10602
## Rationale for this change
This PR adds a ScalarUDF function `to_local_time()`:
- this function converts a timezone-aware timestamp to local time (with no offset or timezone information). In other words, this function strips off the timezone from the timestamp, while keep the display value of the timestamp the same. See examples below
- only accept 1 input with type `Timestamp(..., *)`
- returns with type `Timestamp(..., None)`
### Example
This is how to use it in `datafusion-cli`:
```sql
> select to_local_time('2024-04-01T00:00:20Z'::timestamp AT TIME ZONE 'Europe/Brussels');
+---------------------------------------------+
| to_local_time(Utf8("2024-04-01T00:00:20Z")) |
+---------------------------------------------+
| 2024-04-01T00:00:20 |
+---------------------------------------------+
1 row(s) fetched.
Elapsed 0.010 seconds.
> select to_local_time('2024-04-01T00:00:20'::timestamp AT TIME ZONE 'Europe/Brussels');
+--------------------------------------------+
| to_local_time(Utf8("2024-04-01T00:00:20")) |
+--------------------------------------------+
| 2024-04-01T00:00:20 |
+--------------------------------------------+
1 row(s) fetched.
Elapsed 0.008 seconds.
> select
time,
arrow_typeof(time) as type,
to_local_time(time) as to_local_time,
arrow_typeof(to_local_time(time)) as to_local_time_type
from (
select '2024-04-01T00:00:20Z'::timestamp AT TIME ZONE 'Europe/Brussels' as time
);
+---------------------------+------------------------------------------------+---------------------+-----------------------------+
| time | type | to_local_time | to_local_time_type |
+---------------------------+------------------------------------------------+---------------------+-----------------------------+
| 2024-04-01T00:00:20+02:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) | 2024-04-01T00:00:20 | Timestamp(Nanosecond, None) |
+---------------------------+------------------------------------------------+---------------------+-----------------------------+
1 row(s) fetched.
Elapsed 0.017 seconds.
```
### Example of using `to_local_time()` in `date_bin()`
Combine `to_local_time()` with `date_bin()` will look like:
```sql
> select date_bin(interval '1 day', to_local_time('2024-04-01T00:00:20Z'::timestamp AT TIME ZONE 'Europe/Brussels'));
+----------------------------------------------------------------------------------------------------+
| date_bin(IntervalMonthDayNano("18446744073709551616"),to_local_time(Utf8("2024-04-01T00:00:20Z"))) |
+----------------------------------------------------------------------------------------------------+
| 2024-04-01T00:00:00 |
+----------------------------------------------------------------------------------------------------+
> select date_bin(interval '1 day', to_local_time('2024-04-01T00:00:20Z'::timestamp AT TIME ZONE 'Europe/Brussels')) AT TIME ZONE 'Europe/Brussels';
+----------------------------------------------------------------------------------------------------+
| date_bin(IntervalMonthDayNano("18446744073709551616"),to_local_time(Utf8("2024-04-01T00:00:20Z"))) |
+----------------------------------------------------------------------------------------------------+
| 2024-04-01T00:00:00+02:00 |
+----------------------------------------------------------------------------------------------------+
```
<details>
<summary>Click to see more examples of applying to array values</summary>
<br>
1. Write sample data
```sql
create or replace table t AS
VALUES
('2024-01-01T00:00:01Z'),
('2024-02-01T00:00:01Z'),
('2024-03-01T00:00:01Z'),
('2024-04-01T00:00:01Z'),
('2024-05-01T00:00:01Z'),
('2024-06-01T00:00:01Z'),
('2024-07-01T00:00:01Z'),
('2024-08-01T00:00:01Z'),
('2024-09-01T00:00:01Z'),
('2024-10-01T00:00:01Z'),
('2024-11-01T00:00:01Z'),
('2024-12-01T00:00:01Z')
;
create or replace view t_utc as
select column1::timestamp AT TIME ZONE 'UTC' as "column1"
from t;
create or replace view t_timezone
as
select column1::timestamp AT TIME ZONE 'Europe/Brussels' as "column1"
from t;
```
2. See how tables look like
```sql
> select column1, arrow_typeof(column1) from t;
+----------------------+-------------------------+
| column1 | arrow_typeof(t.column1) |
+----------------------+-------------------------+
| 2024-01-01T00:00:01Z | Utf8 |
| 2024-02-01T00:00:01Z | Utf8 |
| 2024-03-01T00:00:01Z | Utf8 |
| 2024-04-01T00:00:01Z | Utf8 |
| 2024-05-01T00:00:01Z | Utf8 |
| 2024-06-01T00:00:01Z | Utf8 |
| 2024-07-01T00:00:01Z | Utf8 |
| 2024-08-01T00:00:01Z | Utf8 |
| 2024-09-01T00:00:01Z | Utf8 |
| 2024-10-01T00:00:01Z | Utf8 |
| 2024-11-01T00:00:01Z | Utf8 |
| 2024-12-01T00:00:01Z | Utf8 |
+----------------------+-------------------------+
12 row(s) fetched.
Elapsed 0.009 seconds.
> select column1, arrow_typeof(column1) from t_utc;
+----------------------+------------------------------------+
| column1 | arrow_typeof(t_utc.column1) |
+----------------------+------------------------------------+
| 2024-01-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-02-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-03-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-04-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-05-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-06-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-07-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-08-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-09-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-10-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-11-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
| 2024-12-01T00:00:01Z | Timestamp(Nanosecond, Some("UTC")) |
+----------------------+------------------------------------+
12 row(s) fetched.
Elapsed 0.011 seconds.
> select column1, arrow_typeof(column1) from t_timezone;
+---------------------------+------------------------------------------------+
| column1 | arrow_typeof(t_timezone.column1) |
+---------------------------+------------------------------------------------+
| 2024-01-01T00:00:01+01:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-02-01T00:00:01+01:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-03-01T00:00:01+01:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-04-01T00:00:01+02:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-05-01T00:00:01+02:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-06-01T00:00:01+02:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-07-01T00:00:01+02:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-08-01T00:00:01+02:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-09-01T00:00:01+02:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-10-01T00:00:01+02:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-11-01T00:00:01+01:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
| 2024-12-01T00:00:01+01:00 | Timestamp(Nanosecond, Some("Europe/Brussels")) |
+---------------------------+------------------------------------------------+
12 row(s) fetched.
Elapsed 0.012 seconds.
```
3. Query using `to_local_time()`
```sql
> select column1, to_local_time(column1), arrow_typeof(to_local_time(column1)) from t_utc;
+----------------------+------------------------------+--------------------------------------------+
| column1 | to_local_time(t_utc.column1) | arrow_typeof(to_local_time(t_utc.column1)) |
+----------------------+------------------------------+--------------------------------------------+
| 2024-01-01T00:00:01Z | 2024-01-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-02-01T00:00:01Z | 2024-02-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-03-01T00:00:01Z | 2024-03-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-04-01T00:00:01Z | 2024-04-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-05-01T00:00:01Z | 2024-05-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-06-01T00:00:01Z | 2024-06-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-07-01T00:00:01Z | 2024-07-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-08-01T00:00:01Z | 2024-08-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-09-01T00:00:01Z | 2024-09-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-10-01T00:00:01Z | 2024-10-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-11-01T00:00:01Z | 2024-11-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-12-01T00:00:01Z | 2024-12-01T00:00:01 | Timestamp(Nanosecond, None) |
+----------------------+------------------------------+--------------------------------------------+
12 row(s) fetched.
Elapsed 0.015 seconds.
> select column1, to_local_time(column1), arrow_typeof(to_local_time(column1)) from t_timezone;
+---------------------------+-----------------------------------+-------------------------------------------------+
| column1 | to_local_time(t_timezone.column1) | arrow_typeof(to_local_time(t_timezone.column1)) |
+---------------------------+-----------------------------------+-------------------------------------------------+
| 2024-01-01T00:00:01+01:00 | 2024-01-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-02-01T00:00:01+01:00 | 2024-02-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-03-01T00:00:01+01:00 | 2024-03-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-04-01T00:00:01+02:00 | 2024-04-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-05-01T00:00:01+02:00 | 2024-05-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-06-01T00:00:01+02:00 | 2024-06-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-07-01T00:00:01+02:00 | 2024-07-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-08-01T00:00:01+02:00 | 2024-08-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-09-01T00:00:01+02:00 | 2024-09-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-10-01T00:00:01+02:00 | 2024-10-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-11-01T00:00:01+01:00 | 2024-11-01T00:00:01 | Timestamp(Nanosecond, None) |
| 2024-12-01T00:00:01+01:00 | 2024-12-01T00:00:01 | Timestamp(Nanosecond, None) |
+---------------------------+-----------------------------------+-------------------------------------------------+
12 row(s) fetched.
Elapsed 0.016 seconds.
```
4. Combine with `date_bin()`
```sql
> select date_bin(interval '1 day', to_local_time(column1)) AT TIME ZONE 'Europe/Brussels' as date_bin from t_utc;
+---------------------------+
| date_bin |
+---------------------------+
| 2024-01-01T00:00:00+01:00 |
| 2024-02-01T00:00:00+01:00 |
| 2024-03-01T00:00:00+01:00 |
| 2024-04-01T00:00:00+02:00 |
| 2024-05-01T00:00:00+02:00 |
| 2024-06-01T00:00:00+02:00 |
| 2024-07-01T00:00:00+02:00 |
| 2024-08-01T00:00:00+02:00 |
| 2024-09-01T00:00:00+02:00 |
| 2024-10-01T00:00:00+02:00 |
| 2024-11-01T00:00:00+01:00 |
| 2024-12-01T00:00:00+01:00 |
+---------------------------+
12 row(s) fetched.
Elapsed 0.023 seconds.
> select date_bin(interval '1 day', to_local_time(column1)) AT TIME ZONE 'Europe/Brussels' as date_bin from t_timezone;
+---------------------------+
| date_bin |
+---------------------------+
| 2024-01-01T00:00:00+01:00 |
| 2024-02-01T00:00:00+01:00 |
| 2024-03-01T00:00:00+01:00 |
| 2024-04-01T00:00:00+02:00 |
| 2024-05-01T00:00:00+02:00 |
| 2024-06-01T00:00:00+02:00 |
| 2024-07-01T00:00:00+02:00 |
| 2024-08-01T00:00:00+02:00 |
| 2024-09-01T00:00:00+02:00 |
| 2024-10-01T00:00:00+02:00 |
| 2024-11-01T00:00:00+01:00 |
| 2024-12-01T00:00:00+01:00 |
+---------------------------+
12 row(s) fetched.
Elapsed 0.011 seconds.
```
</details>
## What changes are included in this PR?
New ScalarUDF function `to_local_time()` with tests
## Are these changes tested?
Yes
## Are there any user-facing changes?
No API changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671321236
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
Review Comment:
Updated in https://github.com/apache/datafusion/pull/11347/commits/afe619f057eb4c7830f1fa01933527d21fa30d61
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), None),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), Some("+01:00".into())),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampNanosecond(Some(1), Some("Europe/Brussels".into())),
+ )]);
+ assert!(res.is_ok());
+
+ //
+ // Fallible test cases
+ //
+
+ // invalid number of arguments -- no argument
+ let res = ToLocalTimeFunc::new().invoke(&[]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 0"
+ );
+
+ // invalid number of arguments -- more than 1 argument
+ let res = ToLocalTimeFunc::new().invoke(&[
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(Some(1), None)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some("some string".to_string()))),
+ ]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 2"
+ );
+
+ // invalid argument data type
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::Utf8(Some("some string".to_string())),
+ )]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires timestamp argument, got Utf8"
+ );
+
+ // TODO chunchun: invalid timestamp: us, ms, s
Review Comment:
Updated in https://github.com/apache/datafusion/pull/11347/commits/afe619f057eb4c7830f1fa01933527d21fa30d61
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1670518147
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
Review Comment:
Yes, please!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671316167
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
Review Comment:
Good point to mention daylight saving here! Updated in https://github.com/apache/datafusion/pull/11347/commits/ae9f1ad38668d3f080b59f15b7653828861c3bae
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1672265164
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{Array, ArrayRef, PrimitiveBuilder};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse()?;
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T: ArrowTimestampType>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue> {
+ let mut builder = PrimitiveBuilder::<T>::new();
+
+ let primitive_array = as_primitive_array::<T>(array)?;
+ for ts_opt in primitive_array.iter() {
+ match ts_opt {
+ None => builder.append_null(),
+ Some(ts) => {
+ let adjusted_ts: i64 =
+ adjust_to_local_time::<T>(ts, tz)?;
+ builder.append_value(adjusted_ts)
+ }
+ }
+ }
+
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
Review Comment:
Yeah, it would be neat to use `try_unary`. I got the same compiling error when I used `try_unary`, so I rewrote the code in a for loop with `PrimitiveBuilder`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#issuecomment-2223350188
Thank you everyone! This is a great addition i think
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#issuecomment-2218756554
I took the liberty of merging up from main to get the CI error fix
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671321760
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), None),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), Some("+01:00".into())),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampNanosecond(Some(1), Some("Europe/Brussels".into())),
+ )]);
+ assert!(res.is_ok());
+
+ //
+ // Fallible test cases
+ //
+
+ // invalid number of arguments -- no argument
+ let res = ToLocalTimeFunc::new().invoke(&[]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 0"
+ );
+
+ // invalid number of arguments -- more than 1 argument
+ let res = ToLocalTimeFunc::new().invoke(&[
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(Some(1), None)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some("some string".to_string()))),
+ ]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 2"
+ );
+
+ // invalid argument data type
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::Utf8(Some("some string".to_string())),
+ )]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires timestamp argument, got Utf8"
+ );
+
+ // TODO chunchun: invalid timestamp: us, ms, s
+ }
+
+ #[test]
+ fn test_to_local_time_scalar() {
+ let timezone = Some("Europe/Brussels".into());
+ let timestamps_with_timezone = vec![
+ (
+ ScalarValue::TimestampNanosecond(
+ Some(1_123_123_000_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampNanosecond(Some(1_123_130_200_000_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMicrosecond(
+ Some(1_123_123_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMicrosecond(Some(1_123_130_200_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMillisecond(
+ Some(1_123_123_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMillisecond(Some(1_123_130_200_000), None),
+ ),
+ (
+ ScalarValue::TimestampSecond(Some(1_123_123_000), timezone),
+ ScalarValue::TimestampSecond(Some(1_123_130_200), None),
+ ),
+ ];
+
+ for (input, expected) in timestamps_with_timezone {
+ test_to_local_time_helpfer(input, expected);
+ }
+ }
+
+ #[test]
+ fn test_timezone_with_daylight_savings() {
+ let timezone_str = "America/New_York";
+
+ let test_cases = vec![
+ // DST time
+ "2019-03-31T04:00:00",
Review Comment:
Updated in https://github.com/apache/datafusion/pull/11347/commits/db5b73ebb44ee80dcef938dd631395e69b03af1b
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671326748
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
Review Comment:
Thank you for the suggestion! I added error handling in https://github.com/apache/datafusion/pull/11347/commits/afe619f057eb4c7830f1fa01933527d21fa30d61. As a result, this part of the code got rewritten.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "jayzhan211 (via GitHub)" <gi...@apache.org>.
jayzhan211 commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1669748108
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
Review Comment:
I think we can add the tests to `datafusion/sqllogictest/test_files/functions.slt`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671314991
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
Review Comment:
Good call! Updated in https://github.com/apache/datafusion/pull/11347/commits/5280f062c94a17772cf9833c7457742f8e42de45
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#issuecomment-2220519822
@alamb Thank you for the careful review!
> The other thing we should do is document this function in the function reference: [datafusion.apache.org/user-guide/sql/scalar_functions.html](https://datafusion.apache.org/user-guide/sql/scalar_functions.html#time-and-date-functions)
>
> However, we can do that as a follow on PR as this one is already quite large
I'll have a follow-up PR to document this new function
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671319814
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
Review Comment:
I did found these example involves lots of ceremony to prepare the data type. So I move them to the unit test for `adjust_to_local_time()`. I found adding tests in `slt` for `to_local_time()` is more straightforward
Updated in https://github.com/apache/datafusion/pull/11347/commits/043c9971a4ab9106b5d3002ad8ee6b84a76967f0
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1672271610
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{Array, ArrayRef, PrimitiveBuilder};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse()?;
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T: ArrowTimestampType>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue> {
+ let mut builder = PrimitiveBuilder::<T>::new();
+
+ let primitive_array = as_primitive_array::<T>(array)?;
+ for ts_opt in primitive_array.iter() {
+ match ts_opt {
+ None => builder.append_null(),
+ Some(ts) => {
+ let adjusted_ts: i64 =
+ adjust_to_local_time::<T>(ts, tz)?;
+ builder.append_value(adjusted_ts)
+ }
+ }
+ }
+
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ Timestamp(Nanosecond, Some(_)) => {
+ transform_array::<TimestampNanosecondType>(array, tz)
+ }
+ Timestamp(Microsecond, Some(_)) => {
+ transform_array::<TimestampMicrosecondType>(array, tz)
+ }
+ Timestamp(Millisecond, Some(_)) => {
+ transform_array::<TimestampMillisecondType>(array, tz)
+ }
+ Timestamp(Second, Some(_)) => {
+ transform_array::<TimestampSecondType>(array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels" for this date.
+///
+/// Note that the offset varies with daylight savings time (DST), which makes this tricky! For
+/// example, timezone "Europe/Brussels" has a 2-hour offset during DST and a 1-hour offset
+/// when DST ends.
+///
+/// Consequently, DataFusion can represent the timestamp in local time (with no offset or
+/// timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+///
+/// See `test_adjust_to_local_time()` for example
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
+ fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
+ where
+ F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
+ {
+ match converter(ts) {
+ MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
+ "Ambiguous timestamp in microseconds. Do you mean {:?} or {:?}",
+ earliest,
+ latest
+ ),
+ MappedLocalTime::None => exec_err!(
+ "The local time does not exist because there is a gap in the local time."
+ ),
+ MappedLocalTime::Single(date_time) => Ok(date_time),
+ }
+ }
+
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => convert_timestamp(ts, |ts| Utc.timestamp_micros(ts))?,
+ Millisecond => convert_timestamp(ts, |ts| Utc.timestamp_millis_opt(ts))?,
+ Second => convert_timestamp(ts, |ts| Utc.timestamp_opt(ts, 0))?,
+ };
+
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time = date_time.add(
+ // This should not fail under normal circumstances as the
+ // maximum possible offset is 26 hours (93,600 seconds)
+ TimeDelta::try_seconds(offset_seconds)
+ .ok_or(DataFusionError::Internal("Offset seconds should be less than i64::MAX / 1_000 or greater than -i64::MAX / 1_000".to_string()))?,
+ );
+
+ // convert the navie datetime back to i64
Review Comment:
Updated in https://github.com/apache/datafusion/pull/11347/commits/2c35025de8ea30f45264714c9036643d5c494eb2
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1672115527
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{Array, ArrayRef, PrimitiveBuilder};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse()?;
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T: ArrowTimestampType>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue> {
+ let mut builder = PrimitiveBuilder::<T>::new();
+
+ let primitive_array = as_primitive_array::<T>(array)?;
+ for ts_opt in primitive_array.iter() {
+ match ts_opt {
+ None => builder.append_null(),
+ Some(ts) => {
+ let adjusted_ts: i64 =
+ adjust_to_local_time::<T>(ts, tz)?;
+ builder.append_value(adjusted_ts)
+ }
+ }
+ }
+
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
Review Comment:
You could also use [`try_unary`](https://docs.rs/arrow/latest/arrow/compute/fn.try_unary.html) here (that basically does the same thing as what you have here)
```rust
let primitive_array = as_primitive_array::<T>(array)?;
let ts_array = try_unary(primitive_array, |ts| {
adjust_to_local_time::<T>(ts, tz)
})?;
Ok(ColumnarValue::Array(Arc::new(ts_array)))
```
I tried it locally, and using try_unary does require that `adjust_to_local_time` is changed to return ArrowError rather than DataFusion error
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{Array, ArrayRef, PrimitiveBuilder};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse()?;
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T: ArrowTimestampType>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue> {
+ let mut builder = PrimitiveBuilder::<T>::new();
+
+ let primitive_array = as_primitive_array::<T>(array)?;
+ for ts_opt in primitive_array.iter() {
+ match ts_opt {
+ None => builder.append_null(),
+ Some(ts) => {
+ let adjusted_ts: i64 =
+ adjust_to_local_time::<T>(ts, tz)?;
+ builder.append_value(adjusted_ts)
+ }
+ }
+ }
+
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ Timestamp(Nanosecond, Some(_)) => {
+ transform_array::<TimestampNanosecondType>(array, tz)
+ }
+ Timestamp(Microsecond, Some(_)) => {
+ transform_array::<TimestampMicrosecondType>(array, tz)
+ }
+ Timestamp(Millisecond, Some(_)) => {
+ transform_array::<TimestampMillisecondType>(array, tz)
+ }
+ Timestamp(Second, Some(_)) => {
+ transform_array::<TimestampSecondType>(array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels" for this date.
+///
+/// Note that the offset varies with daylight savings time (DST), which makes this tricky! For
+/// example, timezone "Europe/Brussels" has a 2-hour offset during DST and a 1-hour offset
+/// when DST ends.
+///
+/// Consequently, DataFusion can represent the timestamp in local time (with no offset or
+/// timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+///
+/// See `test_adjust_to_local_time()` for example
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
+ fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
+ where
+ F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
+ {
+ match converter(ts) {
+ MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
+ "Ambiguous timestamp in microseconds. Do you mean {:?} or {:?}",
Review Comment:
❤️
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{Array, ArrayRef, PrimitiveBuilder};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse()?;
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T: ArrowTimestampType>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue> {
+ let mut builder = PrimitiveBuilder::<T>::new();
+
+ let primitive_array = as_primitive_array::<T>(array)?;
+ for ts_opt in primitive_array.iter() {
+ match ts_opt {
+ None => builder.append_null(),
+ Some(ts) => {
+ let adjusted_ts: i64 =
+ adjust_to_local_time::<T>(ts, tz)?;
+ builder.append_value(adjusted_ts)
+ }
+ }
+ }
+
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ Timestamp(Nanosecond, Some(_)) => {
+ transform_array::<TimestampNanosecondType>(array, tz)
+ }
+ Timestamp(Microsecond, Some(_)) => {
+ transform_array::<TimestampMicrosecondType>(array, tz)
+ }
+ Timestamp(Millisecond, Some(_)) => {
+ transform_array::<TimestampMillisecondType>(array, tz)
+ }
+ Timestamp(Second, Some(_)) => {
+ transform_array::<TimestampSecondType>(array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels" for this date.
+///
+/// Note that the offset varies with daylight savings time (DST), which makes this tricky! For
+/// example, timezone "Europe/Brussels" has a 2-hour offset during DST and a 1-hour offset
+/// when DST ends.
+///
+/// Consequently, DataFusion can represent the timestamp in local time (with no offset or
+/// timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+///
+/// See `test_adjust_to_local_time()` for example
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
+ fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
+ where
+ F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
+ {
+ match converter(ts) {
+ MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
+ "Ambiguous timestamp in microseconds. Do you mean {:?} or {:?}",
+ earliest,
+ latest
+ ),
+ MappedLocalTime::None => exec_err!(
+ "The local time does not exist because there is a gap in the local time."
+ ),
+ MappedLocalTime::Single(date_time) => Ok(date_time),
+ }
+ }
+
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => convert_timestamp(ts, |ts| Utc.timestamp_micros(ts))?,
+ Millisecond => convert_timestamp(ts, |ts| Utc.timestamp_millis_opt(ts))?,
+ Second => convert_timestamp(ts, |ts| Utc.timestamp_opt(ts, 0))?,
+ };
+
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time = date_time.add(
+ // This should not fail under normal circumstances as the
+ // maximum possible offset is 26 hours (93,600 seconds)
+ TimeDelta::try_seconds(offset_seconds)
+ .ok_or(DataFusionError::Internal("Offset seconds should be less than i64::MAX / 1_000 or greater than -i64::MAX / 1_000".to_string()))?,
+ );
+
+ // convert the navie datetime back to i64
Review Comment:
```suggestion
// convert the naive datetime back to i64
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1672273877
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{Array, ArrayRef, PrimitiveBuilder};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse()?;
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T: ArrowTimestampType>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue> {
+ let mut builder = PrimitiveBuilder::<T>::new();
+
+ let primitive_array = as_primitive_array::<T>(array)?;
+ for ts_opt in primitive_array.iter() {
+ match ts_opt {
+ None => builder.append_null(),
+ Some(ts) => {
+ let adjusted_ts: i64 =
+ adjust_to_local_time::<T>(ts, tz)?;
+ builder.append_value(adjusted_ts)
+ }
+ }
+ }
+
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ Timestamp(Nanosecond, Some(_)) => {
+ transform_array::<TimestampNanosecondType>(array, tz)
+ }
+ Timestamp(Microsecond, Some(_)) => {
+ transform_array::<TimestampMicrosecondType>(array, tz)
+ }
+ Timestamp(Millisecond, Some(_)) => {
+ transform_array::<TimestampMillisecondType>(array, tz)
+ }
+ Timestamp(Second, Some(_)) => {
+ transform_array::<TimestampSecondType>(array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels" for this date.
+///
+/// Note that the offset varies with daylight savings time (DST), which makes this tricky! For
+/// example, timezone "Europe/Brussels" has a 2-hour offset during DST and a 1-hour offset
+/// when DST ends.
+///
+/// Consequently, DataFusion can represent the timestamp in local time (with no offset or
+/// timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+///
+/// See `test_adjust_to_local_time()` for example
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
+ fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
+ where
+ F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
+ {
+ match converter(ts) {
+ MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
+ "Ambiguous timestamp in microseconds. Do you mean {:?} or {:?}",
Review Comment:
I just realized that this is a general function -- the error message applies to microsecond, millisecond, and second. So I removed the phrase `in microseconds` in the error to avoid confusion.
Updated in https://github.com/apache/datafusion/pull/11347/commits/2c35025de8ea30f45264714c9036643d5c494eb2
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{Array, ArrayRef, PrimitiveBuilder};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse()?;
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T: ArrowTimestampType>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue> {
+ let mut builder = PrimitiveBuilder::<T>::new();
+
+ let primitive_array = as_primitive_array::<T>(array)?;
+ for ts_opt in primitive_array.iter() {
+ match ts_opt {
+ None => builder.append_null(),
+ Some(ts) => {
+ let adjusted_ts: i64 =
+ adjust_to_local_time::<T>(ts, tz)?;
+ builder.append_value(adjusted_ts)
+ }
+ }
+ }
+
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ Timestamp(Nanosecond, Some(_)) => {
+ transform_array::<TimestampNanosecondType>(array, tz)
+ }
+ Timestamp(Microsecond, Some(_)) => {
+ transform_array::<TimestampMicrosecondType>(array, tz)
+ }
+ Timestamp(Millisecond, Some(_)) => {
+ transform_array::<TimestampMillisecondType>(array, tz)
+ }
+ Timestamp(Second, Some(_)) => {
+ transform_array::<TimestampSecondType>(array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels" for this date.
+///
+/// Note that the offset varies with daylight savings time (DST), which makes this tricky! For
+/// example, timezone "Europe/Brussels" has a 2-hour offset during DST and a 1-hour offset
+/// when DST ends.
+///
+/// Consequently, DataFusion can represent the timestamp in local time (with no offset or
+/// timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+///
+/// See `test_adjust_to_local_time()` for example
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
+ fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
+ where
+ F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
+ {
+ match converter(ts) {
+ MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
+ "Ambiguous timestamp in microseconds. Do you mean {:?} or {:?}",
Review Comment:
I just realized that this is a general function -- the error message applies to microsecond, millisecond, and second😅 So I removed the phrase `in microseconds` in the error to avoid confusion.
Updated in https://github.com/apache/datafusion/pull/11347/commits/2c35025de8ea30f45264714c9036643d5c494eb2
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#issuecomment-2221254720
> I'll have a follow-up PR to document this new function
Here is the follow-up PR:
- https://github.com/apache/datafusion/pull/11401
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671325483
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
Review Comment:
Thank you for the suggestion! I updated them in two commits:
- https://github.com/apache/datafusion/pull/11347/commits/043c9971a4ab9106b5d3002ad8ee6b84a76967f0 -- parse timezone once
- https://github.com/apache/datafusion/pull/11347/commits/afe619f057eb4c7830f1fa01933527d21fa30d61 -- handle the error rather than unwrap
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671322442
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), None),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), Some("+01:00".into())),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampNanosecond(Some(1), Some("Europe/Brussels".into())),
+ )]);
+ assert!(res.is_ok());
+
+ //
+ // Fallible test cases
+ //
+
+ // invalid number of arguments -- no argument
+ let res = ToLocalTimeFunc::new().invoke(&[]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 0"
+ );
+
+ // invalid number of arguments -- more than 1 argument
+ let res = ToLocalTimeFunc::new().invoke(&[
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(Some(1), None)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some("some string".to_string()))),
+ ]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 2"
+ );
+
+ // invalid argument data type
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::Utf8(Some("some string".to_string())),
+ )]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires timestamp argument, got Utf8"
+ );
+
+ // TODO chunchun: invalid timestamp: us, ms, s
+ }
+
+ #[test]
+ fn test_to_local_time_scalar() {
+ let timezone = Some("Europe/Brussels".into());
+ let timestamps_with_timezone = vec![
+ (
+ ScalarValue::TimestampNanosecond(
+ Some(1_123_123_000_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampNanosecond(Some(1_123_130_200_000_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMicrosecond(
+ Some(1_123_123_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMicrosecond(Some(1_123_130_200_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMillisecond(
+ Some(1_123_123_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMillisecond(Some(1_123_130_200_000), None),
+ ),
+ (
+ ScalarValue::TimestampSecond(Some(1_123_123_000), timezone),
+ ScalarValue::TimestampSecond(Some(1_123_130_200), None),
+ ),
+ ];
+
+ for (input, expected) in timestamps_with_timezone {
+ test_to_local_time_helpfer(input, expected);
+ }
+ }
+
+ #[test]
+ fn test_timezone_with_daylight_savings() {
+ let timezone_str = "America/New_York";
+
+ let test_cases = vec![
+ // DST time
+ "2019-03-31T04:00:00",
+ "2020-03-31T13:40:00",
+ // End of DST
+ "2019-11-04T04:06:40",
+ "2020-11-04T14:06:40",
+ ];
+
+ for timestamp_str in test_cases {
+ // timestamp in local timezone
+ let tz: arrow::array::timezone::Tz =
+ timezone_str.parse().expect("Invalid timezone");
+ let timestamp = timestamp_str
+ .parse::<NaiveDateTime>()
+ .unwrap()
+ .and_local_timezone(tz) // this is in a local timezone
+ .unwrap()
+ .timestamp_nanos_opt()
+ .unwrap();
+
+ // timestamp with no timezone
+ let expected_timestamp = timestamp_str
+ .parse::<NaiveDateTime>()
+ .unwrap()
+ .and_utc() // this is in UTC
+ .timestamp_nanos_opt()
+ .unwrap();
+
+ let input = ScalarValue::TimestampNanosecond(
+ Some(timestamp),
+ Some(timezone_str.into()),
+ );
+ let expected =
+ ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+ test_to_local_time_helpfer(input, expected)
+ }
+ }
+
+ fn test_to_local_time_helpfer(input: ScalarValue, expected: ScalarValue) {
+ let res = ToLocalTimeFunc::new()
+ .invoke(&[ColumnarValue::Scalar(input)])
+ .unwrap();
+ match res {
+ ColumnarValue::Scalar(res) => {
+ assert_eq!(res, expected);
+ }
+ _ => panic!("unexpected return type"),
+ }
+ }
+
+ #[test]
+ fn test_to_local_time_timezones_array() {
+ // TODO chunchun: might need to change string to Timestamp in cases
Review Comment:
This is a leftover comment and no longer needed. I have removed it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671315538
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), None),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), Some("+01:00".into())),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampNanosecond(Some(1), Some("Europe/Brussels".into())),
+ )]);
+ assert!(res.is_ok());
+
+ //
+ // Fallible test cases
+ //
+
+ // invalid number of arguments -- no argument
+ let res = ToLocalTimeFunc::new().invoke(&[]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 0"
+ );
+
+ // invalid number of arguments -- more than 1 argument
+ let res = ToLocalTimeFunc::new().invoke(&[
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(Some(1), None)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some("some string".to_string()))),
+ ]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 2"
+ );
+
+ // invalid argument data type
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::Utf8(Some("some string".to_string())),
+ )]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires timestamp argument, got Utf8"
+ );
+
+ // TODO chunchun: invalid timestamp: us, ms, s
+ }
+
+ #[test]
+ fn test_to_local_time_scalar() {
+ let timezone = Some("Europe/Brussels".into());
+ let timestamps_with_timezone = vec![
+ (
+ ScalarValue::TimestampNanosecond(
+ Some(1_123_123_000_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampNanosecond(Some(1_123_130_200_000_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMicrosecond(
+ Some(1_123_123_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMicrosecond(Some(1_123_130_200_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMillisecond(
+ Some(1_123_123_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMillisecond(Some(1_123_130_200_000), None),
+ ),
+ (
+ ScalarValue::TimestampSecond(Some(1_123_123_000), timezone),
+ ScalarValue::TimestampSecond(Some(1_123_130_200), None),
+ ),
+ ];
+
+ for (input, expected) in timestamps_with_timezone {
+ test_to_local_time_helpfer(input, expected);
+ }
+ }
+
+ #[test]
+ fn test_timezone_with_daylight_savings() {
+ let timezone_str = "America/New_York";
+
+ let test_cases = vec![
+ // DST time
+ "2019-03-31T04:00:00",
+ "2020-03-31T13:40:00",
+ // End of DST
+ "2019-11-04T04:06:40",
+ "2020-11-04T14:06:40",
+ ];
+
+ for timestamp_str in test_cases {
+ // timestamp in local timezone
+ let tz: arrow::array::timezone::Tz =
+ timezone_str.parse().expect("Invalid timezone");
+ let timestamp = timestamp_str
+ .parse::<NaiveDateTime>()
+ .unwrap()
+ .and_local_timezone(tz) // this is in a local timezone
+ .unwrap()
+ .timestamp_nanos_opt()
+ .unwrap();
+
+ // timestamp with no timezone
+ let expected_timestamp = timestamp_str
+ .parse::<NaiveDateTime>()
+ .unwrap()
+ .and_utc() // this is in UTC
+ .timestamp_nanos_opt()
+ .unwrap();
+
+ let input = ScalarValue::TimestampNanosecond(
+ Some(timestamp),
+ Some(timezone_str.into()),
+ );
+ let expected =
+ ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+ test_to_local_time_helpfer(input, expected)
+ }
+ }
+
+ fn test_to_local_time_helpfer(input: ScalarValue, expected: ScalarValue) {
Review Comment:
😅 fixed in https://github.com/apache/datafusion/pull/11347/commits/f8ef51b37aab9de38957a139c306783ef0500099
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1670538440
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
Review Comment:
These examples would be easier to follow if they were only for `adjust_to_local_time` -- there is a lot of setup required to create a `ScalarValue`/`ColumnarValue` and call `invoke`, when I think the idea is really to illustrate how `adjust_to_local_time` works
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
Review Comment:
I don't think it is huge deal as this just clones a bunch of `Arc`s but I think you can avoid this clone like
```suggestion
let time_value = &args[0];
```
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
Review Comment:
This is one of the clearest explanations of what a function does that I have read in a long time 💯 Nice work @appletreeisyellow
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
Review Comment:
is this still a todo?
In theory, if the input was a valid arrow array the value here would be valid too. However, returning an error rather than `panic`'ing is likely nicer
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
Review Comment:
1. I think the timezone parsing should probably also return an error rather than unwrap
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
Review Comment:
TIL https://doc.rust-lang.org/std/vec/struct.Vec.html#method.concat
I think you could potentially make this slightly more efficient using `flat_map` and save collecting an intermediate Vecs:
```suggestion
let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
.into_iter()
.flat_map(base_sig)
.collect::<Vec<_>>()
```
You could then also change the `closure` above to return a `[]` rather than `vec![]`
```rust
[
Exact(vec![Timestamp(array_type, None)]),
Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
]
```
```rust
vec![
Exact(vec![Timestamp(array_type, None)]),
Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
]
```
I don't think this will make a difference performance wise in practice as this struct be created once and then reused many times (via an `Arc<ScalarUDF>`)
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
Review Comment:
likewise here I think it would be good to avoid the panics
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
Review Comment:
It might also be useful to mention here that the offset is not constant, as it varies based on daylight savings time (which makes this so tricky)
Something like
```suggestion
/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
/// offset for "Europe/Brussels" for this date (the offset varies with daylight savings time).
/// Consequently, DataFusion can represent the timestamp in
/// local time (with no offset or timezone information) as
```
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), None),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), Some("+01:00".into())),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampNanosecond(Some(1), Some("Europe/Brussels".into())),
+ )]);
+ assert!(res.is_ok());
+
+ //
+ // Fallible test cases
+ //
+
+ // invalid number of arguments -- no argument
+ let res = ToLocalTimeFunc::new().invoke(&[]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 0"
+ );
+
+ // invalid number of arguments -- more than 1 argument
+ let res = ToLocalTimeFunc::new().invoke(&[
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(Some(1), None)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some("some string".to_string()))),
+ ]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 2"
+ );
+
+ // invalid argument data type
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::Utf8(Some("some string".to_string())),
+ )]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires timestamp argument, got Utf8"
+ );
+
+ // TODO chunchun: invalid timestamp: us, ms, s
+ }
+
+ #[test]
+ fn test_to_local_time_scalar() {
+ let timezone = Some("Europe/Brussels".into());
+ let timestamps_with_timezone = vec![
+ (
+ ScalarValue::TimestampNanosecond(
+ Some(1_123_123_000_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampNanosecond(Some(1_123_130_200_000_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMicrosecond(
+ Some(1_123_123_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMicrosecond(Some(1_123_130_200_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMillisecond(
+ Some(1_123_123_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMillisecond(Some(1_123_130_200_000), None),
+ ),
+ (
+ ScalarValue::TimestampSecond(Some(1_123_123_000), timezone),
+ ScalarValue::TimestampSecond(Some(1_123_130_200), None),
+ ),
+ ];
+
+ for (input, expected) in timestamps_with_timezone {
+ test_to_local_time_helpfer(input, expected);
+ }
+ }
+
+ #[test]
+ fn test_timezone_with_daylight_savings() {
+ let timezone_str = "America/New_York";
+
+ let test_cases = vec![
+ // DST time
+ "2019-03-31T04:00:00",
Review Comment:
I recommend we also hard code the expected timestamp value here (and we can verify it against `timestamp_str.parse()`) -- otherwise this tests is effectively testing that the udf does the same thing as `chrono` rather than get the expected output.
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), None),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), Some("+01:00".into())),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampNanosecond(Some(1), Some("Europe/Brussels".into())),
+ )]);
+ assert!(res.is_ok());
+
+ //
+ // Fallible test cases
+ //
+
+ // invalid number of arguments -- no argument
+ let res = ToLocalTimeFunc::new().invoke(&[]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 0"
+ );
+
+ // invalid number of arguments -- more than 1 argument
+ let res = ToLocalTimeFunc::new().invoke(&[
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(Some(1), None)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some("some string".to_string()))),
+ ]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 2"
+ );
+
+ // invalid argument data type
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::Utf8(Some("some string".to_string())),
+ )]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires timestamp argument, got Utf8"
+ );
+
+ // TODO chunchun: invalid timestamp: us, ms, s
Review Comment:
is it still a TODO?
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), None),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), Some("+01:00".into())),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampNanosecond(Some(1), Some("Europe/Brussels".into())),
+ )]);
+ assert!(res.is_ok());
+
+ //
+ // Fallible test cases
+ //
+
+ // invalid number of arguments -- no argument
+ let res = ToLocalTimeFunc::new().invoke(&[]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 0"
+ );
+
+ // invalid number of arguments -- more than 1 argument
+ let res = ToLocalTimeFunc::new().invoke(&[
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(Some(1), None)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some("some string".to_string()))),
+ ]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 2"
+ );
+
+ // invalid argument data type
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::Utf8(Some("some string".to_string())),
+ )]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires timestamp argument, got Utf8"
+ );
+
+ // TODO chunchun: invalid timestamp: us, ms, s
+ }
+
+ #[test]
+ fn test_to_local_time_scalar() {
+ let timezone = Some("Europe/Brussels".into());
+ let timestamps_with_timezone = vec![
+ (
+ ScalarValue::TimestampNanosecond(
+ Some(1_123_123_000_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampNanosecond(Some(1_123_130_200_000_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMicrosecond(
+ Some(1_123_123_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMicrosecond(Some(1_123_130_200_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMillisecond(
+ Some(1_123_123_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMillisecond(Some(1_123_130_200_000), None),
+ ),
+ (
+ ScalarValue::TimestampSecond(Some(1_123_123_000), timezone),
+ ScalarValue::TimestampSecond(Some(1_123_130_200), None),
+ ),
+ ];
+
+ for (input, expected) in timestamps_with_timezone {
+ test_to_local_time_helpfer(input, expected);
+ }
+ }
+
+ #[test]
+ fn test_timezone_with_daylight_savings() {
+ let timezone_str = "America/New_York";
+
+ let test_cases = vec![
+ // DST time
+ "2019-03-31T04:00:00",
+ "2020-03-31T13:40:00",
+ // End of DST
+ "2019-11-04T04:06:40",
+ "2020-11-04T14:06:40",
+ ];
+
+ for timestamp_str in test_cases {
+ // timestamp in local timezone
+ let tz: arrow::array::timezone::Tz =
+ timezone_str.parse().expect("Invalid timezone");
+ let timestamp = timestamp_str
+ .parse::<NaiveDateTime>()
+ .unwrap()
+ .and_local_timezone(tz) // this is in a local timezone
+ .unwrap()
+ .timestamp_nanos_opt()
+ .unwrap();
+
+ // timestamp with no timezone
+ let expected_timestamp = timestamp_str
+ .parse::<NaiveDateTime>()
+ .unwrap()
+ .and_utc() // this is in UTC
+ .timestamp_nanos_opt()
+ .unwrap();
+
+ let input = ScalarValue::TimestampNanosecond(
+ Some(timestamp),
+ Some(timezone_str.into()),
+ );
+ let expected =
+ ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+ test_to_local_time_helpfer(input, expected)
+ }
+ }
+
+ fn test_to_local_time_helpfer(input: ScalarValue, expected: ScalarValue) {
+ let res = ToLocalTimeFunc::new()
+ .invoke(&[ColumnarValue::Scalar(input)])
+ .unwrap();
+ match res {
+ ColumnarValue::Scalar(res) => {
+ assert_eq!(res, expected);
+ }
+ _ => panic!("unexpected return type"),
+ }
+ }
+
+ #[test]
+ fn test_to_local_time_timezones_array() {
+ // TODO chunchun: might need to change string to Timestamp in cases
Review Comment:
todo?
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), None),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampSecond(Some(1), Some("+01:00".into())),
+ )]);
+ assert!(res.is_ok());
+
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::TimestampNanosecond(Some(1), Some("Europe/Brussels".into())),
+ )]);
+ assert!(res.is_ok());
+
+ //
+ // Fallible test cases
+ //
+
+ // invalid number of arguments -- no argument
+ let res = ToLocalTimeFunc::new().invoke(&[]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 0"
+ );
+
+ // invalid number of arguments -- more than 1 argument
+ let res = ToLocalTimeFunc::new().invoke(&[
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(Some(1), None)),
+ ColumnarValue::Scalar(ScalarValue::Utf8(Some("some string".to_string()))),
+ ]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires 1 argument, got 2"
+ );
+
+ // invalid argument data type
+ let res = ToLocalTimeFunc::new().invoke(&[ColumnarValue::Scalar(
+ ScalarValue::Utf8(Some("some string".to_string())),
+ )]);
+ assert_eq!(
+ res.err().unwrap().strip_backtrace(),
+ "Execution error: to_local_time function requires timestamp argument, got Utf8"
+ );
+
+ // TODO chunchun: invalid timestamp: us, ms, s
+ }
+
+ #[test]
+ fn test_to_local_time_scalar() {
+ let timezone = Some("Europe/Brussels".into());
+ let timestamps_with_timezone = vec![
+ (
+ ScalarValue::TimestampNanosecond(
+ Some(1_123_123_000_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampNanosecond(Some(1_123_130_200_000_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMicrosecond(
+ Some(1_123_123_000_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMicrosecond(Some(1_123_130_200_000_000), None),
+ ),
+ (
+ ScalarValue::TimestampMillisecond(
+ Some(1_123_123_000_000),
+ timezone.clone(),
+ ),
+ ScalarValue::TimestampMillisecond(Some(1_123_130_200_000), None),
+ ),
+ (
+ ScalarValue::TimestampSecond(Some(1_123_123_000), timezone),
+ ScalarValue::TimestampSecond(Some(1_123_130_200), None),
+ ),
+ ];
+
+ for (input, expected) in timestamps_with_timezone {
+ test_to_local_time_helpfer(input, expected);
+ }
+ }
+
+ #[test]
+ fn test_timezone_with_daylight_savings() {
+ let timezone_str = "America/New_York";
+
+ let test_cases = vec![
+ // DST time
+ "2019-03-31T04:00:00",
+ "2020-03-31T13:40:00",
+ // End of DST
+ "2019-11-04T04:06:40",
+ "2020-11-04T14:06:40",
+ ];
+
+ for timestamp_str in test_cases {
+ // timestamp in local timezone
+ let tz: arrow::array::timezone::Tz =
+ timezone_str.parse().expect("Invalid timezone");
+ let timestamp = timestamp_str
+ .parse::<NaiveDateTime>()
+ .unwrap()
+ .and_local_timezone(tz) // this is in a local timezone
+ .unwrap()
+ .timestamp_nanos_opt()
+ .unwrap();
+
+ // timestamp with no timezone
+ let expected_timestamp = timestamp_str
+ .parse::<NaiveDateTime>()
+ .unwrap()
+ .and_utc() // this is in UTC
+ .timestamp_nanos_opt()
+ .unwrap();
+
+ let input = ScalarValue::TimestampNanosecond(
+ Some(timestamp),
+ Some(timezone_str.into()),
+ );
+ let expected =
+ ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+ test_to_local_time_helpfer(input, expected)
+ }
+ }
+
+ fn test_to_local_time_helpfer(input: ScalarValue, expected: ScalarValue) {
Review Comment:
typo?
```suggestion
fn test_to_local_time_helper(input: ScalarValue, expected: ScalarValue) {
```
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
Review Comment:
In general, if possible I think we should use .slt for these tests. The instructions are in https://github.com/apache/datafusion/tree/main/datafusion/sqllogictest
If there is something we can't test easily / completely with .slt tests then I think it is ok to use these unit tests, but it would be good to add a comment about why they aren't tested via .slt
For example the error tests are probably good to put in slt
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
Review Comment:
2. As written this is going to parse the timezone for each row, which will likely be much slower than needed
I think you could do this timezone parsing once in `transform_array` and then pass the parsed timezone into `adjust_to_local_time`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "Abdullahsab3 (via GitHub)" <gi...@apache.org>.
Abdullahsab3 commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1670830134
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
Review Comment:
minor remark: You can have both `array` definitions in the same definition to avoid shadowing
```suggestion
let array = as_primitive_array::<T>(array)?.unary(|ts| adjust_to_local_time::<T>(ts, tz));
```
(my linting might be off :p)
Thinking outloud here, but perhaps it's better to follow the same naming convention as for the scalar types. i.e:
```suggestion
let array = as_primitive_array::<T>(array)?;
let adjusted_array = .unary(|ts| adjust_to_local_time::<T>(ts, tz));
```
in which case the returned value needs to be changed to:
```Rust
Ok(ColumnarValue::Array(Arc::new(adjusted_array)))
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671321056
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
Review Comment:
These are the leftovers I forgot to clean up🙈 Updated in https://github.com/apache/datafusion/pull/11347/commits/afe619f057eb4c7830f1fa01933527d21fa30d61
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#issuecomment-2221081287
I plan to merge this PR tomorrow morning Eastern time unless there are any other comments or anyone would like additional time to review
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671327247
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,610 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse().unwrap(); // TODO chunchun: remove unwrap
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::clone(array)))
Review Comment:
Thank you for the suggestion! Closing `time_value` is more clean! Updated in https://github.com/apache/datafusion/pull/11347/commits/b8be3b1e222f998c75ec77d4c0a41094531d504b
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671322760
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
Review Comment:
Neat! Updated in https://github.com/apache/datafusion/pull/11347/commits/5fe18d5e12c2479bfeaf0ce6cc84cdcf7670c93d
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#issuecomment-2218883245
@alamb @jayzhan211 @Abdullahsab3 -- Thank you for all your reviews and the many helpful suggestions ♥️ I have addressed all of them. I'm happy to make further updates if there are any additional suggestions.
@alamb -- Thank you for merging the clippy error fix!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "Abdullahsab3 (via GitHub)" <gi...@apache.org>.
Abdullahsab3 commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1670852619
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,610 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse().unwrap(); // TODO chunchun: remove unwrap
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::clone(array)))
Review Comment:
Minor remark: Is it needed to create a new enum instance in this case? Wondering whether cloning `time_value` might be better in this case
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
Review Comment:
minor remark: You can have both `array` definitions in the same definition to avoid shadowing
```suggestion
let array = as_primitive_array::<T>(array)?.unary(|ts| adjust_to_local_time::<T>(ts, tz));
```
(my linting might be off :p)
Thinking outloud here, but perhaps it's better to follow the same naming convention as for the scalar types. i.e:
```suggestion
let array = as_primitive_array::<T>(array)?;
let adjusted_array = .unary(|ts| adjust_to_local_time::<T>(ts, tz));
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "Abdullahsab3 (via GitHub)" <gi...@apache.org>.
Abdullahsab3 commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1670830134
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
Review Comment:
minor remark: You can have both `array` definitions in the same definition to avoid shadowing
```suggestion
let array = as_primitive_array::<T>(array)?.unary(|ts| adjust_to_local_time::<T>(ts, tz));
```
(my linting might be off :p)
Thinking outloud here, but perhaps it's better to follow the same naming convention as for the scalar types. i.e:
```suggestion
let array = as_primitive_array::<T>(array)?;
let adjusted_array = array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
```
in which case the returned value needs to be changed to:
```Rust
Ok(ColumnarValue::Array(Arc::new(adjusted_array)))
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1672273877
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,564 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{Array, ArrayRef, PrimitiveBuilder};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{DateTime, MappedLocalTime, Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, DataFusionError, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ [
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .flat_map(base_sig)
+ .collect::<Vec<_>>();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = &args[0];
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // If has timezone, adjust the underlying time value. The current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC. Therefore,
+ // we need to adjust the time value to the local time. See [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(timezone)) => {
+ let tz: Tz = timezone.parse()?;
+
+ match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(_),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(*ts, tz)?;
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T: ArrowTimestampType>(
+ array: &ArrayRef,
+ tz: Tz,
+ ) -> Result<ColumnarValue> {
+ let mut builder = PrimitiveBuilder::<T>::new();
+
+ let primitive_array = as_primitive_array::<T>(array)?;
+ for ts_opt in primitive_array.iter() {
+ match ts_opt {
+ None => builder.append_null(),
+ Some(ts) => {
+ let adjusted_ts: i64 =
+ adjust_to_local_time::<T>(ts, tz)?;
+ builder.append_value(adjusted_ts)
+ }
+ }
+ }
+
+ Ok(ColumnarValue::Array(Arc::new(builder.finish())))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ Timestamp(Nanosecond, Some(_)) => {
+ transform_array::<TimestampNanosecondType>(array, tz)
+ }
+ Timestamp(Microsecond, Some(_)) => {
+ transform_array::<TimestampMicrosecondType>(array, tz)
+ }
+ Timestamp(Millisecond, Some(_)) => {
+ transform_array::<TimestampMillisecondType>(array, tz)
+ }
+ Timestamp(Second, Some(_)) => {
+ transform_array::<TimestampSecondType>(array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels" for this date.
+///
+/// Note that the offset varies with daylight savings time (DST), which makes this tricky! For
+/// example, timezone "Europe/Brussels" has a 2-hour offset during DST and a 1-hour offset
+/// when DST ends.
+///
+/// Consequently, DataFusion can represent the timestamp in local time (with no offset or
+/// timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+///
+/// See `test_adjust_to_local_time()` for example
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, tz: Tz) -> Result<i64> {
+ fn convert_timestamp<F>(ts: i64, converter: F) -> Result<DateTime<Utc>>
+ where
+ F: Fn(i64) -> MappedLocalTime<DateTime<Utc>>,
+ {
+ match converter(ts) {
+ MappedLocalTime::Ambiguous(earliest, latest) => exec_err!(
+ "Ambiguous timestamp in microseconds. Do you mean {:?} or {:?}",
Review Comment:
I just realized that this is a general function -- the error message applies to microsecond, millisecond, and second. So I
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#issuecomment-2218755706
The clippy error was fixed in https://github.com/apache/datafusion/pull/11368
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "appletreeisyellow (via GitHub)" <gi...@apache.org>.
appletreeisyellow commented on code in PR #11347:
URL: https://github.com/apache/datafusion/pull/11347#discussion_r1671314166
##########
datafusion/functions/src/datetime/to_local_time.rs:
##########
@@ -0,0 +1,601 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::any::Any;
+use std::ops::Add;
+use std::sync::Arc;
+
+use arrow::array::timezone::Tz;
+use arrow::array::{ArrayRef, PrimitiveArray};
+use arrow::datatypes::DataType::Timestamp;
+use arrow::datatypes::{
+ ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType,
+ TimestampNanosecondType, TimestampSecondType,
+};
+use arrow::datatypes::{
+ TimeUnit,
+ TimeUnit::{Microsecond, Millisecond, Nanosecond, Second},
+};
+
+use chrono::{Offset, TimeDelta, TimeZone, Utc};
+use datafusion_common::cast::as_primitive_array;
+use datafusion_common::{exec_err, Result, ScalarValue};
+use datafusion_expr::TypeSignature::Exact;
+use datafusion_expr::{
+ ColumnarValue, ScalarUDFImpl, Signature, Volatility, TIMEZONE_WILDCARD,
+};
+
+/// A UDF function that converts a timezone-aware timestamp to local time (with no offset or
+/// timezone information). In other words, this function strips off the timezone from the timestamp,
+/// while keep the display value of the timestamp the same.
+///
+/// # Example 1
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// // 2019-03-31 01:00:00 +01:00
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+/// Some(1_553_990_400),
+/// Some("Europe/Brussels".into()),
+/// ))])
+/// .unwrap();
+///
+/// // 2019-03-31 01:00:00 <-- this timestamp no longer has +01:00 offset
+/// let expected = ScalarValue::TimestampSecond(Some(1_553_994_000), None);
+///
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+///
+/// # Example 2
+///
+/// ```
+/// # use datafusion_common::ScalarValue;
+/// # use datafusion_expr::ColumnarValue;
+/// # use chrono::NaiveDateTime;
+/// # use datafusion_functions::datetime::to_local_time::ToLocalTimeFunc;
+/// # use datafusion_expr::ScalarUDFImpl;
+///
+/// let timestamp_str = "2020-03-31T13:40:00";
+/// let timezone_str = "America/New_York";
+/// let tz: arrow::array::timezone::Tz =
+/// timezone_str.parse().expect("Invalid timezone");
+///
+/// let timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_local_timezone(tz) // this is in a local timezone
+/// .unwrap()
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let expected_timestamp = timestamp_str
+/// .parse::<NaiveDateTime>()
+/// .unwrap()
+/// .and_utc() // this is in UTC
+/// .timestamp_nanos_opt()
+/// .unwrap();
+///
+/// let input =
+/// ScalarValue::TimestampNanosecond(Some(timestamp), Some(timezone_str.into()));
+/// let res = ToLocalTimeFunc::new()
+/// .invoke(&[ColumnarValue::Scalar(input)])
+/// .unwrap();
+/// let expected = ScalarValue::TimestampNanosecond(Some(expected_timestamp), None);
+/// match res {
+/// ColumnarValue::Scalar(res) => {
+/// assert_eq!(res, expected);
+/// }
+/// _ => panic!("unexpected return type"),
+/// }
+/// ```
+#[derive(Debug)]
+pub struct ToLocalTimeFunc {
+ signature: Signature,
+}
+
+impl Default for ToLocalTimeFunc {
+ fn default() -> Self {
+ Self::new()
+ }
+}
+
+impl ToLocalTimeFunc {
+ pub fn new() -> Self {
+ let base_sig = |array_type: TimeUnit| {
+ vec![
+ Exact(vec![Timestamp(array_type, None)]),
+ Exact(vec![Timestamp(array_type, Some(TIMEZONE_WILDCARD.into()))]),
+ ]
+ };
+
+ let full_sig = [Nanosecond, Microsecond, Millisecond, Second]
+ .into_iter()
+ .map(base_sig)
+ .collect::<Vec<_>>()
+ .concat();
+
+ Self {
+ signature: Signature::one_of(full_sig, Volatility::Immutable),
+ }
+ }
+
+ fn to_local_time(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {}",
+ args.len()
+ );
+ }
+
+ let time_value = args[0].clone();
+ let arg_type = time_value.data_type();
+ match arg_type {
+ DataType::Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(time_value.clone())
+ }
+ // if has timezone, adjust the underlying time value. the current time value
+ // is stored as i64 in UTC, even though the timezone may not be in UTC, so
+ // we need to adjust the time value to the local time. see [`adjust_to_local_time`]
+ // for more details.
+ //
+ // Then remove the timezone in return type, i.e. return None
+ DataType::Timestamp(_, Some(_)) => match time_value {
+ ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampNanosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMicrosecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampMillisecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(ts),
+ Some(tz),
+ )) => {
+ let adjusted_ts =
+ adjust_to_local_time::<TimestampSecondType>(ts, &tz);
+ Ok(ColumnarValue::Scalar(ScalarValue::TimestampSecond(
+ Some(adjusted_ts),
+ None,
+ )))
+ }
+ ColumnarValue::Array(array) => {
+ fn transform_array<T>(
+ array: &ArrayRef,
+ tz: &str,
+ ) -> Result<ColumnarValue>
+ where
+ T: ArrowTimestampType,
+ {
+ let array = as_primitive_array::<T>(array)?;
+ let array: PrimitiveArray<T> =
+ array.unary(|ts| adjust_to_local_time::<T>(ts, tz));
+
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+
+ match array.data_type() {
+ Timestamp(_, None) => {
+ // if no timezone specificed, just return the input
+ Ok(ColumnarValue::Array(Arc::new(array)))
+ }
+ Timestamp(Nanosecond, Some(tz)) => {
+ transform_array::<TimestampNanosecondType>(&array, tz)
+ }
+ Timestamp(Microsecond, Some(tz)) => {
+ transform_array::<TimestampMicrosecondType>(&array, tz)
+ }
+ Timestamp(Millisecond, Some(tz)) => {
+ transform_array::<TimestampMillisecondType>(&array, tz)
+ }
+ Timestamp(Second, Some(tz)) => {
+ transform_array::<TimestampSecondType>(&array, tz)
+ }
+ _ => {
+ exec_err!("to_local_time function requires timestamp argument in array, got {:?}", array.data_type())
+ }
+ }
+ }
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ time_value.data_type()
+ )
+ }
+ },
+ _ => {
+ exec_err!(
+ "to_local_time function requires timestamp argument, got {:?}",
+ arg_type
+ )
+ }
+ }
+ }
+}
+
+/// This function converts a timestamp with a timezone to a timestamp without a timezone.
+/// The display value of the adjusted timestamp remain the same, but the underlying timestamp
+/// representation is adjusted according to the relative timezone offset to UTC.
+///
+/// This function uses chrono to handle daylight saving time changes.
+///
+/// For example,
+///
+/// ```text
+/// '2019-03-31T01:00:00Z'::timestamp at time zone 'Europe/Brussels'
+/// ```
+///
+/// is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00+01:00
+/// ```
+///
+/// and is represented in DataFusion as:
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_990_400_000_000_000), Some("Europe/Brussels"))
+/// ```
+///
+/// To strip off the timezone while keeping the display value the same, we need to
+/// adjust the underlying timestamp with the timezone offset value using `adjust_to_local_time()`
+///
+/// ```text
+/// adjust_to_local_time(1_553_990_400_000_000_000, "Europe/Brussels") --> 1_553_994_000_000_000_000
+/// ```
+///
+/// The difference between `1_553_990_400_000_000_000` and `1_553_994_000_000_000_000` is
+/// `3600_000_000_000` ns, which corresponds to 1 hour. This matches with the timezone
+/// offset for "Europe/Brussels". Consequently, DataFusion can represent the timestamp in
+/// local time (with no offset or timezone information) as
+///
+/// ```text
+/// TimestampNanosecond(Some(1_553_994_000_000_000_000), None)
+/// ```
+///
+/// which is displayed as follows in datafusion-cli:
+///
+/// ```text
+/// 2019-03-31T01:00:00
+/// ```
+fn adjust_to_local_time<T: ArrowTimestampType>(ts: i64, timezone: &str) -> i64 {
+ let date_time = match T::UNIT {
+ Nanosecond => Utc.timestamp_nanos(ts),
+ Microsecond => Utc.timestamp_micros(ts).unwrap(), // TODO chunchun: replace unwrap
+ Millisecond => Utc.timestamp_millis_opt(ts).unwrap(),
+ Second => Utc.timestamp_opt(ts, 0).unwrap(),
+ };
+
+ let tz: Tz = timezone.parse().unwrap();
+ let offset_seconds: i64 = tz
+ .offset_from_utc_datetime(&date_time.naive_utc())
+ .fix()
+ .local_minus_utc() as i64;
+
+ let adjusted_date_time =
+ date_time.add(TimeDelta::try_seconds(offset_seconds).unwrap());
+
+ // convert the navie datetime back to i64
+ match T::UNIT {
+ Nanosecond => adjusted_date_time.timestamp_nanos_opt().unwrap(), // TODO chuynchun: remove unwrap
+ Microsecond => adjusted_date_time.timestamp_micros(),
+ Millisecond => adjusted_date_time.timestamp_millis(),
+ Second => adjusted_date_time.timestamp(),
+ }
+}
+
+impl ScalarUDFImpl for ToLocalTimeFunc {
+ fn as_any(&self) -> &dyn Any {
+ self
+ }
+
+ fn name(&self) -> &str {
+ "to_local_time"
+ }
+
+ fn signature(&self) -> &Signature {
+ &self.signature
+ }
+
+ fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
+ if arg_types.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ arg_types.len()
+ );
+ }
+
+ match &arg_types[0] {
+ Timestamp(Nanosecond, _) => Ok(Timestamp(Nanosecond, None)),
+ Timestamp(Microsecond, _) => Ok(Timestamp(Microsecond, None)),
+ Timestamp(Millisecond, _) => Ok(Timestamp(Millisecond, None)),
+ Timestamp(Second, _) => Ok(Timestamp(Second, None)),
+ _ => exec_err!(
+ "The to_local_time function can only accept timestamp as the arg, got {:?}", arg_types[0]
+ ),
+ }
+ }
+
+ fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
+ if args.len() != 1 {
+ return exec_err!(
+ "to_local_time function requires 1 argument, got {:?}",
+ args.len()
+ );
+ }
+
+ self.to_local_time(args)
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use std::sync::Arc;
+
+ use arrow::array::{types::TimestampNanosecondType, TimestampNanosecondArray};
+ use arrow::compute::kernels::cast_utils::string_to_timestamp_nanos;
+ use arrow::datatypes::{DataType, TimeUnit};
+ use chrono::NaiveDateTime;
+ use datafusion_common::ScalarValue;
+ use datafusion_expr::{ColumnarValue, ScalarUDFImpl};
+
+ use super::ToLocalTimeFunc;
+
+ #[test]
+ fn test_to_local_time() {
Review Comment:
@jayzhan211 @alamb Thank you for the suggestions! I found `datafusion/sqllogictest/test_files/timestamps.slt` is more fitting to add the tests for `to_local_time()` function, so I add them there.
Updated in https://github.com/apache/datafusion/pull/11347/commits/3456517854ff0f3f533acc7e9ffcbbc7f4ee7acc
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add UDF to_local_time() [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #11347:
URL: https://github.com/apache/datafusion/pull/11347
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org