Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(rust, python): let "ambiguous" take "null" value #11606

Closed
wants to merge 1 commit into from
Closed

feat(rust, python): let "ambiguous" take "null" value #11606

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 5 additions & 4 deletions crates/polars-arrow/src/kernels/time.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,14 @@ pub fn convert_to_naive_local(
to_tz: &Tz,
ndt: NaiveDateTime,
ambiguous: &str,
) -> Result<NaiveDateTime> {
) -> Result<Option<NaiveDateTime>> {
let ndt = from_tz.from_utc_datetime(&ndt).naive_local();
match to_tz.from_local_datetime(&ndt) {
LocalResult::Single(dt) => Ok(dt.naive_utc()),
LocalResult::Single(dt) => Ok(Some(dt.naive_utc())),
LocalResult::Ambiguous(dt_earliest, dt_latest) => match ambiguous {
"earliest" => Ok(dt_earliest.naive_utc()),
"latest" => Ok(dt_latest.naive_utc()),
"null" => Ok(None),
"earliest" => Ok(Some(dt_earliest.naive_utc())),
"latest" => Ok(Some(dt_latest.naive_utc())),
"raise" => Err(ArrowError::InvalidArgumentError(
format!("datetime '{}' is ambiguous in time zone '{}'. Please use `ambiguous` to tell how it should be localized.", ndt, to_tz)
)),
Expand Down
25 changes: 16 additions & 9 deletions crates/polars-ops/src/chunked_array/datetime/replace_time_zone.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,21 +43,28 @@ pub fn replace_time_zone(
};
let out = match ambiguous.len() {
1 => match unsafe { ambiguous.get_unchecked(0) } {
Some(ambiguous) => datetime.0.try_apply(|timestamp| {
let ndt = timestamp_to_datetime(timestamp);
Ok(datetime_to_timestamp(convert_to_naive_local(
&from_tz, &to_tz, ndt, ambiguous,
)?))
}),
Some(ambiguous) => {
let iter = datetime.0.downcast_iter().map(|arr| {
let element_iter = arr.iter().map(|timestamp_opt| match timestamp_opt {
Some(timestamp) => {
let ndt = timestamp_to_datetime(*timestamp);
let res = convert_to_naive_local(&from_tz, &to_tz, ndt, ambiguous)?;
Ok::<_, PolarsError>(res.map(datetime_to_timestamp))
},
None => Ok(None),
});
element_iter.try_collect_arr()
});
ChunkedArray::try_from_chunk_iter(datetime.0.name(), iter)
},
Comment on lines -46 to +59
Copy link
Collaborator Author

@MarcoGorelli MarcoGorelli Oct 9, 2023 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm I'm seeing a slight perf regression here. small, but looks real?

https://www.kaggle.com/code/marcogorelli/polars-timing?scriptVersionId=145820321

here:

  • min: 34.58047724000062
  • max: 34.82023904700054
  • avg: 34.7417660610001 +/- 0.02811013567654737

main:

  • min: 33.56916854799965
  • max: 34.206052164999164
  • avg: 33.85647739228547 +/- 0.08204153267322174

I don't know, is there a way to do this that might preserve performance?

_ => Ok(datetime.0.apply(|_| None)),
},
_ => try_binary_elementwise(datetime, ambiguous, |timestamp_opt, ambiguous_opt| {
match (timestamp_opt, ambiguous_opt) {
(Some(timestamp), Some(ambiguous)) => {
let ndt = timestamp_to_datetime(timestamp);
Ok(Some(datetime_to_timestamp(convert_to_naive_local(
&from_tz, &to_tz, ndt, ambiguous,
)?)))
let res = convert_to_naive_local(&from_tz, &to_tz, ndt, ambiguous)?;
Ok(res.map(datetime_to_timestamp))
},
_ => Ok(None),
}
Expand Down
3 changes: 3 additions & 0 deletions py-polars/polars/expr/datetime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ def truncate(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Notes
-----
Expand Down Expand Up @@ -264,6 +265,7 @@ def round(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Notes
-----
Expand Down Expand Up @@ -1379,6 +1381,7 @@ def replace_time_zone(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Examples
--------
Expand Down
2 changes: 2 additions & 0 deletions py-polars/polars/expr/string.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ def to_datetime(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Examples
--------
Expand Down Expand Up @@ -252,6 +253,7 @@ def strptime(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Notes
-----
Expand Down
2 changes: 1 addition & 1 deletion py-polars/polars/functions/as_datatype.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ def datetime_(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime

- ``'null'``: set to null

Returns
-------
Expand Down
3 changes: 3 additions & 0 deletions py-polars/polars/series/datetime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1028,6 +1028,7 @@ def replace_time_zone(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Examples
--------
Expand Down Expand Up @@ -1468,6 +1469,7 @@ def truncate(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Notes
-----
Expand Down Expand Up @@ -1667,6 +1669,7 @@ def round(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Returns
-------
Expand Down
2 changes: 2 additions & 0 deletions py-polars/polars/series/string.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ def to_datetime(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Examples
--------
Expand Down Expand Up @@ -235,6 +236,7 @@ def strptime(
- ``'raise'`` (default): raise
- ``'earliest'``: use the earliest datetime
- ``'latest'``: use the latest datetime
- ``'null'``: set to null

Notes
-----
Expand Down
2 changes: 1 addition & 1 deletion py-polars/polars/type_aliases.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,7 @@
] # ListToStructWidthStrategy

# The following have no equivalent on the Rust side
Ambiguous: TypeAlias = Literal["earliest", "latest", "raise"]
Ambiguous: TypeAlias = Literal["earliest", "latest", "raise", "null"]
ConcatMethod = Literal[
"vertical", "vertical_relaxed", "diagonal", "horizontal", "align"
]
Expand Down
29 changes: 29 additions & 0 deletions py-polars/tests/unit/datatypes/test_temporal.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2234,6 +2234,35 @@ def test_replace_time_zone_sortedness_expressions(
assert result["ts"].flags["SORTED_ASC"] == expected_sortedness


def test_replace_time_zone_ambiguous_null() -> None:
df = pl.DataFrame(
{
"a": [datetime(2020, 10, 25, 1)] * 3,
"b": ["earliest", "latest", "null"],
}
)
# expression containing 'null'
result = df.select(
pl.col("a").dt.replace_time_zone("Europe/London", ambiguous=pl.col("b"))
)["a"]
expected = [
datetime(2020, 10, 25, 1, fold=0, tzinfo=ZoneInfo("Europe/London")),
datetime(2020, 10, 25, 1, fold=1, tzinfo=ZoneInfo("Europe/London")),
None,
]
assert result[0] == expected[0]
assert result[1] == expected[1]
assert result[2] == expected[2]

# single 'null' value
result = df.select(
pl.col("a").dt.replace_time_zone("Europe/London", ambiguous="null")
)["a"]
assert result[0] is None
assert result[1] is None
assert result[2] is None


def test_use_earliest_deprecation() -> None:
# strptime
with pytest.warns(
Expand Down