在Polars中根据每行指定的时区检索值的`convert_time_zone`函数

Question

我正在尝试根据每一行指定的时区来确定时间，使用的是 Polars。看看下面这段代码：

df = pl.DataFrame({
    "time": [datetime(2023, 4, 3, 2), datetime(2023, 4, 4, 3), datetime(2023, 4, 5, 4)],
    "tzone": ["Asia/Tokyo", "America/Chicago", "Europe/Paris"]
}).with_columns(c.time.dt.replace_time_zone("UTC"))

df.with_columns(
    tokyo=c.time.dt.convert_time_zone("Asia/Tokyo").dt.hour(),
    chicago=c.time.dt.convert_time_zone("America/Chicago").dt.hour(),
    paris=c.time.dt.convert_time_zone("Europe/Paris").dt.hour()
)

在这个例子中，我为每个时区单独计算了时间，以达到想要的结果，也就是 [11, 22, 6]，这些数字对应于 time 列中根据 tzone 时区的小时数。即便如此，从正确的列中收集信息还是很困难。

不幸的是，下面这个简单的尝试，想要直接将 tzone 列中的时区动态传递给 convert_time_zone 函数，却没有成功：

df.with_columns(c.time.dt.convert_time_zone(c.tzone).dt.hour())
#TypeError: argument 'time_zone': 'Expr' object cannot be converted to 'PyString'

那么，完成这个任务的最优雅的方法是什么呢？

Answer 1

还有一种方法可以做到这一点，而不需要额外的插件，那就是通过时区来拆分源数据，这样你就可以把它提取为一个独特的值。

如果你是从急切模式开始的，可以像这样使用 partition_by：

pl.concat([
    df_inner.lazy().with_columns(
        hour=pl.col('time')
        .dt.replace_time_zone('UTC')
        .dt.convert_time_zone(x[0])
        .dt.hour())
        for x,df_inner in df.with_row_index('__i').partition_by(['tzone'], as_dict=True).items()
]).sort('__i').drop('__i').collect()

这样做会让内部框架变得懒惰，这样在收集数据时，concat 就可以并行处理所有的时区操作。

如果你是从懒惰模式开始的，你可以这样做：

(
    df.lazy()
    .with_row_index('__i')
    .group_by('tzone')
    .map_groups(lambda df : (
        df.with_columns(
            hour=pl.col('time')
                .dt.replace_time_zone('UTC')
                .dt.convert_time_zone(df['tzone'][0])
                .dt.hour())
        ), schema={'time':pl.Datetime, 'tzone':pl.String, 'hour':pl.Int8, '__i':pl.UInt32})
    .sort('__i')
    .drop('__i')
    .collect()
    )

不过这种方法比第一种慢很多，因为它没有并行处理分组。

Answer 2

正如@Hericks在评论中提到的，你可以使用多个when/then条件：

df.with_columns(
   pl.when(c.tzone == tzone)
     .then(c.time.dt.convert_time_zone(tzone).dt.hour())
     .alias(tzone.rsplit("/", 1)[1].lower())
   for tzone in df.get_column("tzone").unique()
)

shape: (3, 5)
┌─────────────────────────┬─────────────────┬───────┬─────────┬───────┐
│ time                    ┆ tzone           ┆ tokyo ┆ chicago ┆ paris │
│ ---                     ┆ ---             ┆ ---   ┆ ---     ┆ ---   │
│ datetime[μs, UTC]       ┆ str             ┆ i8    ┆ i8      ┆ i8    │
╞═════════════════════════╪═════════════════╪═══════╪═════════╪═══════╡
│ 2023-04-03 02:00:00 UTC ┆ Asia/Tokyo      ┆ 11    ┆ null    ┆ null  │
│ 2023-04-04 03:00:00 UTC ┆ America/Chicago ┆ null  ┆ 22      ┆ null  │
│ 2023-04-05 04:00:00 UTC ┆ Europe/Paris    ┆ null  ┆ null    ┆ 6     │
└─────────────────────────┴─────────────────┴───────┴─────────┴───────┘

pl.coalesce()可以用来创建一个单独的列。

df.with_columns(
    pl.coalesce(
       pl.when(c.tzone == tzone)
         .then(c.time.dt.convert_time_zone(tzone).dt.hour())
       for tzone in df.get_column("tzone").unique()
    )
    .alias("hour")
)

shape: (3, 3)
┌─────────────────────────┬─────────────────┬──────┐
│ time                    ┆ tzone           ┆ hour │
│ ---                     ┆ ---             ┆ ---  │
│ datetime[μs, UTC]       ┆ str             ┆ i8   │
╞═════════════════════════╪═════════════════╪══════╡
│ 2023-04-03 02:00:00 UTC ┆ Asia/Tokyo      ┆ 11   │
│ 2023-04-04 03:00:00 UTC ┆ America/Chicago ┆ 22   │
│ 2023-04-05 04:00:00 UTC ┆ Europe/Paris    ┆ 6    │
└─────────────────────────┴─────────────────┴──────┘

Answer 3

要完全实现这个功能，并且支持懒执行（也就是按需执行），唯一的方法是使用 polars-xdt 插件：

df = pl.DataFrame(
    {
        "time": [
            datetime(2023, 4, 3, 2),
            datetime(2023, 4, 4, 3),
            datetime(2023, 4, 5, 4),
        ],
        "tzone": ["Asia/Tokyo", "America/Chicago", "Europe/Paris"],
    }
).with_columns(pl.col("time").dt.replace_time_zone("UTC"))

df.with_columns(
    result=xdt.to_local_datetime("time", pl.col("tzone")).dt.hour(),
)

结果：

Out[6]:
shape: (3, 3)
┌─────────────────────────┬─────────────────┬────────┐
│ time                    ┆ tzone           ┆ result │
│ ---                     ┆ ---             ┆ ---    │
│ datetime[μs, UTC]       ┆ str             ┆ i8     │
╞═════════════════════════╪═════════════════╪════════╡
│ 2023-04-03 02:00:00 UTC ┆ Asia/Tokyo      ┆ 11     │
│ 2023-04-04 03:00:00 UTC ┆ America/Chicago ┆ 22     │
│ 2023-04-05 04:00:00 UTC ┆ Europe/Paris    ┆ 6      │
└─────────────────────────┴─────────────────┴────────┘

https://github.com/pola-rs/polars-xdt

如果你不需要懒执行，那么就像其他回答提到的，你可以遍历你 'time_zone' 列中的唯一值。