如何将Scanner投影列直接用于新数据集Hive分区？

加载Parquet数据集并定义投影后，尝试用重命名后的列作为分区字段保存时出现错误：

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

# 加载原始数据集
myds = ds.dataset('mypath', format='parquet', partitioning='hive')
# 原始schema（注释为计划的重命名）
# On/Off_Peak: string
# area: string
# price: decimal128(8, 4)
# date: date32[day]
# hourbegin: int32
# hourend: int32
# inflation: string rename to Inflation
# Price_Type: string
# Reference_Year: int32
# Case: string
# region: string rename to Region

# 定义投影（包含列重命名和计算字段）
projection = {
    'Region': ds.field('region'),
    'Date': ds.field('date'),
    'isPeak': pc.equal(ds.field('On/Off_Peak'), ds.scalar('On')),
    'Hourbegin': ds.field('hourbegin'),
    'Hourend': ds.field('hourend'),
    'Inflation': ds.field('inflation'),
    'Price_Type': ds.field('Price_Type'),
    'Area': ds.field('area'),
    'Price': ds.field('price'),
    'Reference_Year': ds.field('Reference_Year'),
    'Case': ds.field('Case'),
}

# 创建扫描器
scanner = myds.scanner(columns=projection)

# 尝试保存分区数据集时出错
ds.write_dataset(scanner, 'newpath',
                 partitioning=['Reference_Year', 'Case', 'Region'],
                 partitioning_flavor='hive',
                 format='parquet')

错误信息：

KeyError: 'Column Region does not exist in schema'

当前可行的规避方法是使用原始列名region作为分区字段，之后手动修改目录名，但希望找到直接实现的方法；同时询问如果分区需要更复杂的计算，是否必须先保存非分区数据集再重新分区。

1. 直接使用重命名后的列作为分区字段 #

问题根源是write_dataset直接接收scanner时，会优先参考原始数据集的schema而非投影后的schema。有两种直接解决的方式：

方式一：转换为Table后写入

先将scanner的结果转换为pyarrow.Table，此时Table的schema已包含投影后的列名，可直接用于分区：

# 将扫描器结果转换为Table
table = scanner.to_table()

# 写入分区数据集，直接使用重命名后的列名
ds.write_dataset(table, 'newpath',
                 partitioning=['Reference_Year', 'Case', 'Region'],
                 partitioning_flavor='hive',
                 format='parquet')

方式二：使用字段表达式定义分区

在partitioning参数中直接用字段表达式指定重命名逻辑，无需转换Table：

ds.write_dataset(scanner, 'newpath',
                 partitioning=[
                     ds.field('Reference_Year'),
                     ds.field('Case'),
                     ds.field('region').alias('Region')  # 直接在分区定义中完成重命名
                 ],
                 partitioning_flavor='hive',
                 format='parquet')

2. 复杂计算字段作为分区的处理 #

不需要先保存非分区数据集再重新分区。可以直接在投影中定义计算后的字段，然后将其作为分区字段，用上述两种方式实现：

例如，使用计算字段isPeak作为分区：

# 定义包含计算字段的投影
projection = {
    # 其他字段...
    'isPeak': pc.equal(ds.field('On/Off_Peak'), ds.scalar('On')),
}

scanner = myds.scanner(columns=projection)
table = scanner.to_table()

# 直接使用计算后的字段作为分区
ds.write_dataset(table, 'newpath',
                 partitioning=['Reference_Year', 'Case', 'Region', 'isPeak'],
                 partitioning_flavor='hive',
                 format='parquet')

或者直接在分区参数中指定计算逻辑：

ds.write_dataset(scanner, 'newpath',
                 partitioning=[
                     ds.field('Reference_Year'),
                     ds.field('Case'),
                     ds.field('region').alias('Region'),
                     pc.equal(ds.field('On/Off_Peak'), ds.scalar('On')).alias('isPeak')
                 ],
                 partitioning_flavor='hive',
                 format='parquet')

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