基于PyArrow处理混合时区数据的日期特征计算问题

你的核心问题在于PyArrow的timestamp列仅支持单一时区，无法像Pandas那样为每个时间戳单独绑定时区。结合你数据的特性（同一series_id对应统一时区），这里有两种可行的处理方案，其中分组处理是最优选择：

方案一：按series_id分组处理（推荐） #

因为你的数据中同一series_id的时区完全一致，利用PyArrow的分组API可以高效完成时区转换和特征计算，避免逐行处理的性能损耗：

代码示例 #

import pyarrow as pa
import pyarrow.compute as pc

# 模拟PostgreSQL中的数据结构：UTC时间戳+关联时区
data = pa.table({
    "series_id": [1, 1, 2, 2, 3, 3],
    "timestamp_utc": pa.array([
        "2023-03-06T14:00:00", "2023-03-07T14:00:00",
        "2023-03-06T13:00:00", "2023-03-07T13:00:00",
        "2023-03-06T13:30:00", "2023-03-07T13:30:00"
    ], type=pa.timestamp('us')),
    "timezone": ["Australia/Brisbane", "Australia/Brisbane",
                 "Australia/Sydney", "Australia/Sydney",
                 "Australia/Adelaide", "Australia/Adelaide"]
})

# 定义分组处理逻辑：将UTC时间转成本地时区并计算日期特征
def process_group(group):
    # 提取组内统一的时区
    tz = group["timezone"][0].as_py()
    # 为UTC时间添加时区信息并转换为本地时区的timestamp类型
    local_timestamp = pc.cast(pc.add_timezone(group["timestamp_utc"], tz), pa.timestamp('us', tz=tz))
    # 计算本地时间的日、工作日/周末等特征
    return group.append_columns([
        ("local_day", pc.day(local_timestamp)),
        ("is_weekend", pc.is_weekend(local_timestamp))
    ])

# 执行分组处理
result_table = data.group_by("series_id").apply(process_group)
print(result_table.to_pandas())

这种方式完全契合PyArrow的设计规范，分组后每组的timestamp列可以使用统一时区，计算结果准确且性能优异。

方案二：逐行自定义处理（仅适用于特殊场景） #

如果存在同一series_id对应多个时区的极端情况，可以使用map_elements实现逐行时区转换，但性能会明显低于分组处理：

代码示例 #

def convert_and_get_day(utc_ts, tz):
    from datetime import datetime
    import pytz
    utc_dt = utc_ts.as_py()
    local_dt = utc_dt.astimezone(pytz.timezone(tz.as_py()))
    return local_dt.day

# 注册自定义函数
get_local_day = pc.register_scalar_function(
    convert_and_get_day,
    args=[pa.timestamp('us'), pa.string()],
    return_type=pa.int32()
)

# 计算本地日
local_day_col = get_local_day(data["timestamp_utc"], data["timezone"])
result_table = data.append_column("local_day", local_day_col)

为什么你之前的转换会出错？ #

当你将带多时区的Pandas DataFrame转成PyArrow表时，PyArrow会自动将所有时间戳转换为第一个元素的时区对应的UTC值，并将整个列的时区设为第一个元素的时区。此时调用pa.compute.day()时，会基于列的时区解析UTC时间，导致后续时间戳的日期计算错误。

内容的提问来源于stack exchange，提问作者David Waterworth