feat(python): support DataFrame init with Datetime dtypes that specify a timezone

[alexander-beedie]

Can now initialize DataFrame with timezone-aware DataType; previously had to post-convert per-column.

Example

from datetime import datetime
import polars as pl

df = pl.DataFrame(
    data={
        "d1": [datetime(2022,10,12,12,30)],
        "d2": [datetime(2022,10,12,12,30)],
    },
    columns=[
        ("d1", pl.Datetime(time_zone="America/New_York")),
        ("d2", pl.Datetime(time_zone="Asia/Tokyo")),
    ],
)

Before:

# NotImplementedError: Conversion of polars data type 
#  datetime[μs, America/New_York] to Python type not implemented.

After:

# ┌────────────────────────────────┬──────────────────────────┐
# │ d1                             ┆ d2                       │
# │ ---                            ┆ ---                      │
# │ datetime[μs, America/New_York] ┆ datetime[μs, Asia/Tokyo] │
# ╞════════════════════════════════╪══════════════════════════╡
# │ 2022-10-12 12:30:00 EDT        ┆ 2022-10-12 12:30:00 JST  │
# └────────────────────────────────┴──────────────────────────┘

df.row(0)

# (datetime(2022,10,12, 8,30,tzinfo=ZoneInfo(key='America/New_York')),
#  datetime(2022,10,12,21,30,tzinfo=ZoneInfo(key='Asia/Tokyo')))

Misc:

• Also improves/simplifies some type datatype lookups/conversions.

[ritchie46]

Hmm.. This has some strange locale issues. :/

On my laptop (in the Netherlands) it fails with:

   def test_init_with_timezone() -> None:
        for tu in DTYPE_TEMPORAL_UNITS | frozenset([None]):
            df = pl.DataFrame(
                data={
                    "d1": [datetime(2022, 10, 12, 12, 30)],
                    "d2": [datetime(2022, 10, 12, 12, 30)],
                },
                columns=[
                    ("d1", pl.Datetime(tu, "America/New_York")),  # type: ignore[arg-type]
                    ("d2", pl.Datetime(tu, "Asia/Tokyo")),  # type: ignore[arg-type]
                ],
            )
            # note: setting timezone doesn't change the underlying/physical value...
            assert (df["d1"].to_physical() == df["d2"].to_physical()).all()
    
            # ...but (as expected) it _does_ change the interpretation of that value
>           assert df.rows() == [
                (
                    datetime(2022, 10, 12, 8, 30, tzinfo=ZoneInfo("America/New_York")),
                    datetime(2022, 10, 12, 21, 30, tzinfo=ZoneInfo("Asia/Tokyo")),
                )
            ]
E           AssertionError: assert [(datetime.da...sia/Tokyo')))] == [(datetime.da...sia/Tokyo')))]
E             At index 0 diff: (datetime.datetime(2022, 10, 12, 7, 30, tzinfo=zoneinfo.ZoneInfo(key='America/New_York')), datetime.datetime(2022, 10, 12, 20, 30, tzinfo=zoneinfo.ZoneInfo(key='Asia/Tokyo'))) != (datetime.datetime(2022, 10, 12, 8, 30, tzinfo=zoneinfo.ZoneInfo(key='America/New_York')), datetime.datetime(2022, 10, 12, 21, 30, tzinfo=zoneinfo.ZoneInfo(key='Asia/Tokyo')))
E             Use -v to get more diff

tests/unit/test_df.py:2414: AssertionError

[alexander-beedie]

@ritchie46 argh... must be some kind of implicit local/UTC shenanigans on load into pyarrow from naive datetime objects - what's your current UTC offset? (+1 or +2?) I'll try and harden the test ASAP.

[ritchie46]

@ritchie46 argh... must be some kind of implicit local/UTC shenanigans on load into pyarrow from naive datetime objects - what's your current UTC offset? (+1 or +2?) I'll try and harden the test ASAP.

I am not entirely sure, but I think UTC +2
https://www.timeanddate.com/time/zone/netherlands/amsterdam

[alexander-beedie]

Cheers; think I have a probable fix... one sec.

[alexander-beedie]

Fingers crossed: #5177 🙏