使用Awswrangler时partition_filter报KeyError：source字段实际存在

我尝试从AWS中的Parquet文件加载pd.DataFrame，希望筛选出source列等于crawler_name变量的数据集。使用wr.s3.read_parquet时配置了自定义的partition_filter，但执行时出现KeyError: 'source'；不使用该过滤器时，能正常获取包含对应crawler_name值的完整DataFrame。

相关代码：

def read_aws_parquet(aws_uri: str, crawler_name: str, year: str, month: str, day:Optional[str], columns: list = ['url', 'crawl_timestamp', 'source']) -> pd.DataFrame:

    my_filter = lambda x: True if x["source"] == crawler_name else False
    # my_filter = lambda x: x["source"] == crawler_name
    logger.warning(f'{aws_uri} is a valid url')
    return wr.s3.read_parquet(
        path=aws_uri,
        path_suffix=".snappy.parquet",
        use_threads=True,
        columns=columns,  # 'deactivated_timestamp' must be added later
        partition_filter=my_filter,
        dataset=True
        # pyarrow_additional_kwargs={"filters": [('source', '=', crawler_name)]}
    )

报错信息：

File "/Users/administrator/Library/Caches/pypoetry/virtualenvs/pf-11294-manual-review-of-the-tracking-url-MC2_y5HC-py3.10/lib/python3.10/site-packages/awswrangler/s3/_read_parquet.py", line 749, in read_parquet
    paths = _apply_partition_filter(path_root=path_root, paths=paths, filter_func=partition_filter)
  File "/Users/administrator/Library/Caches/pypoetry/virtualenvs/pf-11294-manual-review-of-the-tracking-url-MC2_y5HC-py3.10/lib/python3.10/site-packages/awswrangler/s3/_read.py", line 88, in _apply_partition_filter
    return [p for p in paths if filter_func(_extract_partitions_from_path(path_root=path_root, path=p)) is True]
  File "/Users/administrator/Library/Caches/pypoetry/virtualenvs/pf-11294-manual-review-of-the-tracking-url-MC2_y5HC-py3.10/lib/python3.10/site-packages/awswrangler/s3/_read.py", line 88, in <listcomp>
    return [p for p in paths if filter_func(_extract_partitions_from_path(path_root=path_root, path=p)) is True]
  File "<input>", line 19, in <lambda>
KeyError: 'source'

原因与解决方案 #

核心原因 #

partition_filter的作用是过滤S3路径中的分区字段（比如路径形如s3://bucket/path/source=crawler_name/year=2024/中的source、year），它会从文件路径里提取分区键值对，而不是读取Parquet文件内部的列。你的source是Parquet文件内部的列，不是路径中的分区字段，所以_extract_partitions_from_path无法从路径里提取到source键，导致抛出KeyError。

解决方案 #

根据你的实际情况选择以下两种方式：

1. 如果source是文件内部列（非分区字段）

直接弃用partition_filter，改用pyarrow_additional_kwargs中的filters参数来过滤文件内部的列，这个参数会在读取Parquet文件时直接筛选数据：

def read_aws_parquet(aws_uri: str, crawler_name: str, year: str, month: str, day:Optional[str], columns: list = ['url', 'crawl_timestamp', 'source']) -> pd.DataFrame:
    logger.warning(f'{aws_uri} is a valid url')
    return wr.s3.read_parquet(
        path=aws_uri,
        path_suffix=".snappy.parquet",
        use_threads=True,
        columns=columns,
        dataset=True,
        pyarrow_additional_kwargs={"filters": [('source', '=', crawler_name)]}
    )

或者也可以在读取完整DataFrame后，用Pandas的筛选逻辑处理：

df = wr.s3.read_parquet(...)
df = df[df['source'] == crawler_name]

2. 如果source应该是分区字段但路径未配置

检查你的S3数据存储路径是否采用Hive风格的分区格式（例如s3://bucket/path/source=your_crawler/year=2024/month=05/），如果不是，需要：

• 重新组织S3上的数据，按分区字段创建对应的目录结构
• 或者更新对应的Glue表分区定义，让AWS Wrangler能识别路径中的source分区字段

完成后，partition_filter就能正常从路径中提取source键进行筛选了。

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