Polars SQL Context日期过滤查询失效问题求助

问题描述 #

使用Polars SQL Context查询Parquet数据时，通过pyarrow的read_table预先传入日期过滤条件能正常返回结果，但直接在SQL查询语句中写日期过滤会报错；尝试用cast转换日期字符串后虽不报错，但无数据返回。

有效代码示例 #

filters = define_filters(event.get("filters", None))
table = pq.read_table(
    f"s3://my_s3_path{partition_path}",
    partitioning="hive",
    filters=filters,
)
df = pl.from_arrow(table)
ctx = pl.SQLContext(stuff=df)
sql = "SELECT things FROM stuff"
new_df = ctx.execute(sql,eager=True)

无效代码示例（filters==None时） #

filters = define_filters(event.get("filters", None))
table = pq.read_table(
    f"s3://my_s3_path{partition_path}",
    partitioning="hive",
    filters=filters,
)
df = pl.from_arrow(table)
ctx = pl.SQLContext(stuff=df)
sql = """
    SELECT things 
    FROM stuff 
    where START_DATE_KEY >= '2023-06-01' and START_DATE_KEY < '2023-06-17'
"""
new_df = ctx.execute(sql,eager=True)

报错信息 #

Traceback (most recent call last):
  File "/Users/xaras/projects/arrow-lambda/loose.py", line 320, in <module>
    test_runner(target=args.t, limit=args.limit, is_debug=args.debug)
  File "/Users/xaras/projects/arrow-lambda/loose.py", line 294, in test_runner
    rows, metadata = test_handler(target, limit, display=True, is_debug=is_debug)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/xaras/projects/arrow-lambda/loose.py", line 110, in test_handler
    response = json.loads(handler(event, None))
                          ^^^^^^^^^^^^^^^^^^^^
  File "/Users/xaras/projects/arrow-lambda/serverless/app.py", line 42, in handler
    new_df = ctx.execute(
             ^^^^^^^^^^^^
  File "/Users/xaras/.pyenv/versions/pyarrow/lib/python3.11/site-packages/polars/sql/context.py", line 275, in execute
    return res.collect() if (eager or self._eager_execution) else res
           ^^^^^^^^^^^^^
  File "/Users/xaras/.pyenv/versions/pyarrow/lib/python3.11/site-packages/polars/utils/deprecation.py", line 95, in wrapper
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/xaras/.pyenv/versions/pyarrow/lib/python3.11/site-packages/polars/lazyframe/frame.py", line 1713, in collect
    return wrap_df(ldf.collect())
                   ^^^^^^^^^^^^^
exceptions.ComputeError: cannot compare 'date/datetime/time' to a string value (create native python { 'date', 'datetime', 'time' } or compare to a temporal column)

补充信息 #

尝试使用cast('2023-06-01' as date)进行日期转换，代码可正常运行，但返回结果为空。

复现示例 #

import polars as pl

df = pl.DataFrame(
    {
        "a": ["2023-06-01", "2023-06-01", "2023-06-02"],
        "b": [None, None, None],
        "c": [4, 5, 6],
        "d": [None, None, None],
    }
)
df = df.with_columns(pl.col("a").str.strptime(pl.Date, "%Y-%m-%d", strict=False))
print(df)
ctx = pl.SQLContext(stuff=df)
new_df = ctx.execute(
    "select * from stuff where a = cast('2023-06-01' as date)",
    eager=True,
)
print(new_df)

运行结果

shape: (3, 4)
┌────────────┬──────┬─────┬──────┐
│ a          ┆ b    ┆ c   ┆ d    │
│ ---        ┆ ---  ┆ --- ┆ ---  │
│ date       ┆ f32  ┆ i64 ┆ f32  │
╞════════════╪══════╪═════╪══════╡
│ 2023-06-01 ┆ null ┆ 4   ┆ null │
│ 2023-06-01 ┆ null ┆ 5   ┆ null │
│ 2023-06-02 ┆ null ┆ 6   ┆ null │
└────────────┴──────┴─────┴──────┘
shape: (0, 4)
┌──────┬─────┬─────┬─────┐
│ a    ┆ b   ┆ c   ┆ d   │
│ ---  ┆ --- ┆ --- ┆ --- │
│ date ┆ f32 ┆ i64 ┆ f32 │
╞══════╪═════╪═════╪═════╡
└──────┴─────┴─────┴─────┘

解决方法 #

Polars SQL当前对cast('YYYY-MM-DD' as date)的解析存在兼容问题，可通过以下两种方式解决：

1. 使用Polars标准日期字面量语法 #

在SQL中直接用date 'YYYY-MM-DD'格式声明日期常量，这是Polars SQL原生支持的日期写法，能正确匹配date类型列：

# 等值查询
new_df = ctx.execute(
    "select * from stuff where a = date '2023-06-01'",
    eager=True,
)

# 范围查询
sql = """
    SELECT things 
    FROM stuff 
    where START_DATE_KEY >= date '2023-06-01' and START_DATE_KEY < date '2023-06-17'
"""
new_df = ctx.execute(sql, eager=True)

2. 绑定Python原生日期对象 #

将日期转换为Python的datetime.date对象，通过params参数传入SQL语句，完全避免字符串解析问题：

from datetime import date

# 等值查询
target_date = date(2023, 6, 1)
new_df = ctx.execute(
    "select * from stuff where a = ?",
    params=[target_date],
    eager=True,
)

# 范围查询
start_date = date(2023, 6, 1)
end_date = date(2023, 6, 17)
new_df = ctx.execute(
    "SELECT things FROM stuff where START_DATE_KEY >= ? and START_DATE_KEY < ?",
    params=[start_date, end_date],
    eager=True,
)

原理说明 #

报错提示“无法将日期类型与字符串比较”，是因为直接写字符串会被当作文本值，无法和Polars的date类型列匹配。而cast('2023-06-01' as date)的转换逻辑存在兼容问题，导致转换后的日期值与列中存储的日期无法匹配。使用date 'YYYY-MM-DD'字面量或绑定原生日期对象，能让Polars正确识别日期类型，实现准确的比较过滤。

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