在 pyarrow 数据集中整合片段的最佳方法？

Question

下面的代码有没有更好的方法来实现我的目标：

1. 我想从包含很多小片段的数据集 A 中读取基于过滤器的数据（这个数据集中的文件很多，因为我经常下载数据）
2. 我想在循环中根据分区合并片段（当我无法将所有过滤器都放入内存时，我正在使用它，因此我一个一个地处理它们）
3. 我想将此数据写入一个新数据集（数据集 B）到一个合并文件中，该文件由我们的 BI 工具读取 - 遗憾的是没有 partition_filename_cb 函数，因此我需要为此使用旧的 write_to_dataset - 这个文件通常是分区名称
4. 我真的很想清理 Dataset A。随着时间的推移，越来越多的文件被添加到分区中，因为我经常下载数据并且可以更新行（其中一些片段文件只有 1 或 2 条记录）李>

以下是我目前的流程。 我使用 ds.Scanner 应用我的过滤器并从原始数据集中选择我的列

def retrieve_fragments(dataset, filter_expression, columns):
    """Creates a dictionary of file fragments and filters from a pyarrow dataset"""
    fragment_partitions = {}
    scanner = ds.Scanner.from_dataset(dataset, columns=columns, filter=filter_expression)
    fragments = scanner.get_fragments()
    for frag in fragments:
        keys = ds._get_partition_keys(frag.partition_expression)
        fragment_partitions[frag] = keys
    return fragment_partitions

下面我创建了所有具有相同过滤器表达式的片段的小列表。然后我可以将这些新数据集写入一个合并文件中，我假设我也可以删除单个片段文件并编写一个新的合并版本？

fragments = retrieve_fragments(
    dataset=dataset, filter_expression=filter_expression, columns=read_columns
)
unique_filters = []
dfs = []
for fragment, filter_value in fragments.items():
    if filter_value not in unique_filters:
        unique_filters.append(filter_value)


#each chunk is a list of all of the fragments with the same partition_expression / filter which we turn into a new dataset that we can then process or resave into a consolidated file
for unique_filter in unique_filters:
    chunks = []
    for frag, filter_value in fragments.items():
        if filter_value == unique_filter:
            chunks.append(frag.path)
    logging.info(
        f"Combining {len(chunks)} fragments with filter {unique_filter} into a single table"
    )
    table = ds.dataset(chunks, partitioning=partitioning, filesystem=filesystem).to_table(columns=read_columns)


    #ignoring metadata due to some issues with columns having a boolean type even though they were never boolean
    df = table.to_pandas(ignore_metadata=True)
    #this function would just sort and drop duplicates on a unique constraint key
    df = prepare_dataframe(df)
    table = pa.Table.from_pandas(df=df, schema=dataset.schema, preserve_index=False)

    #write dataset to Dataset B (using partition_filename_cb)
    #I believe I could now also write the dataset back to Dataset A in a consolidated parquet file and then delete all of the fragment.paths. This would leave me with only a single file in the partition "folder"

此操作的输出会将每个分区的单个文件保存到新数据集中 (/dev/interactions-final/created_date=2019-11-13/2019-11-13.parquet)

INFO - Combining 78 fragments with filter {'created_date': datetime.date(2019, 11, 13)} into a single table
INFO - Saving 172657 rows and 36 columns (70.36 MB to dev/interactions-final)
INFO - Combining 57 fragments with filter {'created_date': datetime.date(2019, 11, 18)} into a single table
INFO - Saving 67036 rows and 36 columns (29.63 MB to dev/interactions-final)
INFO - Combining 55 fragments with filter {'created_date': datetime.date(2019, 11, 19)} into a single table
INFO - Saving 65035 rows and 36 columns (29.62 MB to dev/interactions-final)
INFO - Combining 63 fragments with filter {'created_date': datetime.date(2019, 11, 20)} into a single table
INFO - Saving 63613 rows and 36 columns (30.76 MB to dev/interactions-final)

Answer 1

你试过write_dataset（代码here）吗？它会重新分区，我认为它会在此过程中收集小碎片。