pyarrow - 如何将扫描仪中的投影列合并到新的数据集分区中

Question

假设我加载了一个数据集

myds=ds.dataset('mypath', format='parquet', partitioning='hive')
myds.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'

我可以通过将 partitioning 更改为 ['Reference_Year', 'Case', 'region'] 来解决这个问题，以匹配非投影列(然后稍后更改所有这些目录的名称)但有没有办法直接执行此操作？

假设我的分区需要计算的不仅仅是更改列名。我是否必须在一个步骤中保存非分区数据集以获取新列，然后执行另一次保存操作以创建分区数据集？

Answer 1

编辑:此错误已在 pyarrow 10.0.0 中修复

在我看来这像是一个错误。就好像 write_dataset 正在查看 dataset_schema而不是 projected_schema

我想你可以通过调用 to_reader 来绕过它在扫描仪上。

table = pa.Table.from_arrays(
    [
        pa.array(['a', 'b', 'c'], pa.string()),
        pa.array(['a', 'b', 'c'], pa.string()),
    ],
    names=['region', "Other"]
)
table_dataset = ds.dataset(table)
columns={
    "Region": ds.field('region'),
    "Other": ds.field('Other'),
}
scanner = table_dataset.scanner(columns=columns)

ds.write_dataset(
    scanner.to_reader(), 
    'newpath',
    partitioning=['Region'], partitioning_flavor='hive',
    format='parquet')

我已报告问题 here