如何在PyArrow中基于字典序构造日期范围过滤表达式？

我有一个包含year、month、day整数列的数据集，给定start_date和end_date（pandas Timestamp类型），需要构造过滤表达式筛选出该日期范围内的所有行。

现有两种可行实现 #

实现1：手动按年>月>日优先级构造表达式 #

import pandas as pd
import pyarrow.dataset as ds
import pyarrow.compute as pc

def get_partition_filter(start_time: pd.Timestamp,
                         end_time: pd.Timestamp) -> pc.Expression:
    # 提取起止日期的年、月、日
    start_year, start_month, start_day = start_time.year, start_time.month, start_time.day
    end_year, end_month, end_day = end_time.year, end_time.month, end_time.day

    # 构造过滤表达式
    return ((ds.field("year") > start_year) | ((ds.field("year") == start_year) & (ds.field("month") > start_month)) |
           ((ds.field("year") == start_year) & (ds.field("month") == start_month) & (ds.field("day") >= start_day))) & \
           ((ds.field("year") < end_year) | ((ds.field("year") == end_year) & (ds.field("month") < end_month)) |
           ((ds.field("year") == end_year) & (ds.field("month") == end_month) & (ds.field("day") <= end_day)))

实现2：转为整数派生字段规避字典序判断 #

def get_partition_filter_new(start_time: pd.Timestamp,
                             end_time: pd.Timestamp) -> pc.Expression:
    # 将日期转为整数格式（YYYYMMDD）
    start_date_int, end_date_int = (
        int(time.strftime('%Y%m%d')) for time in (start_time, end_time)
    )

    # 构造派生字段：year*10000 + month*100 + day
    derived_field = (
        ds.field("year")*10**4 +
        ds.field("month")*10**2 +
        ds.field("day")
    )

    # 构造过滤表达式
    return (derived_field >= start_date_int) & (derived_field <= end_date_int)

需求：原生实现元组范围比较逻辑 #

我想知道PyArrow中是否有原生方式实现如下伪代码的逻辑：

# PSEUDO-CODE
(start_year_int, start_month_int, start_day_int) <= (ds.field("year"), ds.field("month"), ds.field("day")) <= (end_year_int, end_month_int, end_day_int) 

解决方案：PyArrow原生结构体比较 #

PyArrow支持结构体（struct）的字典序比较，完全可以实现你想要的元组范围过滤逻辑，代码更简洁且与伪代码逻辑一致：

def get_partition_filter_tuple(start_time: pd.Timestamp, end_time: pd.Timestamp) -> pc.Expression:
    start_year, start_month, start_day = start_time.year, start_time.month, start_time.day
    end_year, end_month, end_day = end_time.year, end_time.month, end_time.day
    
    # 将year、month、day字段组合成结构体
    date_struct = pc.make_struct(ds.field("year"), ds.field("month"), ds.field("day"))
    # 构造起止日期的结构体标量
    start_struct = pc.scalar((start_year, start_month, start_day), type=date_struct.type)
    end_struct = pc.scalar((end_year, end_month, end_day), type=date_struct.type)
    
    # 直接进行范围比较，逻辑与伪代码完全匹配
    return (date_struct >= start_struct) & (date_struct <= end_struct)

说明 #

• PyArrow的结构体比较遵循字典序，和元组比较逻辑完全一致：先比较第一个元素，相等则比较第二个，以此类推。
• 这种方式无需手动拼接复杂的多条件逻辑，也不需要计算派生字段，代码可读性更高，同时是PyArrow原生支持的操作，性能有保障。
• 也可以用pc.struct({"year": ds.field("year"), "month": ds.field("month"), "day": ds.field("day")})的方式构造结构体，效果相同。

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