Dask pyarrow-识别写入拼花地板数据集时使用的片段或过滤器？

我的用例是，我希望将文件路径或过滤器作为xcom传递给Airflow中的任务，以便我的下一个任务可以读取刚刚处理的数据

任务A将表写入分区数据集，并生成大量拼花文件片段-->任务B稍后将这些片段作为数据集读取。不过，我只需要读取相关数据，而不需要读取可能有数百万行的整个数据集

我测试了两种方法：

在我完成写入数据集后立即列出修改过的文件。这将为我提供一个路径列表，我可以在下一个任务中调用ds.dataset（路径）。我可以在这些路径上使用

partitioning.parse（）

，或者检查片段以获得所用过滤器的列表（

frag.partition\u expression

）

这样做的一个缺陷是，我可以将文件并行写入同一个数据集

我可以通过将表转换为数据帧、执行groupby，然后构造过滤器来生成写入数据集时使用的过滤器。我不确定是否有更简单的方法来解决这个问题。然后，我可以对结果使用pq.\u filters\u to\u expression（）来创建可用的过滤器

这并不理想，因为我需要修复某些数据类型，这些数据类型不能正确保存为xcom（没有酸洗，所以所有内容都必须是json格式）。另外，如果我想在字典列上进行分区，我可能需要调整这个函数

    def create_filter_list(df, partition_columns):
        """Creates a list of pyarrow filters to be sent through an xcom and evaluated as an expression. Xcom disables pickling, so we need to save timestamp and date values as strings and convert downstream"""
        filter_list = []
        value_list = []
        partition_keys = [df[col] for col in partition_columns]
        for keys, _ in df[partition_columns].groupby(partition_keys):
            if len(partition_columns) == 1:
                if is_jsonable(keys):
                    value_list.append(keys)
                elif keys is not None:
                    value_list.append(str(keys))
            else:
                if not isinstance(keys, tuple):
                    keys = (keys,)
                read_filter = []
                for name, val in zip(partition_columns, keys):
                    if type(val) == np.int_:
                        read_filter.append((name, "==", int(val)))
                    elif val is not None:
                        read_filter.append((name, "==", str(val)))
                filter_list.append(read_filter)
        if len(partition_columns) == 1:
            if len(value_list) > 0:
                filter_list = [(name, "in", value_list) for name in partition_columns]
        return filter_list

关于我应该采取哪种方法，或者是否有更好的方法来实现我的目标，有什么建议吗？

建议（不确定这是否适合您的用例）：

关键问题是需要正确选择数据的子集，这可以在上游“修复”。更新big dataframe的函数/脚本可以包含一个条件，用于保存已修改的临时数据副本，并在单独（临时）路径中满足某些要求。然后，该文件将被传递给下游任务，这些任务可以在处理临时文件后将其删除。

您可以查看此问题（），我相信它会满足您的需要。同时，您可以使用

basename\u模板

作为解决方法

import glob
import os

import pyarrow as pa
import pyarrow.dataset as pads

class TrackingWriter:

    def __init__(self):
        self.counter = 0
        part_schema = pa.schema({'part': pa.int64()})
        self.partitioning = pads.HivePartitioning(part_schema)

    def next_counter(self):
        result = self.counter
        self.counter += 1
        return result

    def write_dataset(self, table, base_dir):
        counter = self.next_counter()
        pads.write_dataset(table, base_dir, format='parquet', partitioning=self.partitioning, basename_template=f'batch-{counter}-part-{{i}}')
        files_written = glob.glob(os.path.join(base_dir, '**', f'batch-{counter}-*'))
        return files_written

table_one = pa.table({'part': [0, 0, 1, 1], 'val': [1, 2, 3, 4]})
table_two = pa.table({'part': [0, 0, 1, 1], 'val': [5, 6, 7, 8]})

writer = TrackingWriter()
print(writer.write_dataset(table_one, '/tmp/mydataset'))
print(writer.write_dataset(table_two, '/tmp/mydataset'))

---

这只是一个粗略的草图。您可能还希望代码在启动时运行，以查看

计数器的下一个自由值是多少。或者你可以使用
uuid
而不是计数器。
Nice，所以我用adlfs替换操作系统，用aiffair中的“{task\u id}-{execution}”替换“batch-{counter}”，它根据唯一的时间表命名每个文件（这很好，因为我可以并行地将数据加载到同一个数据集中，并且可以使用glob查找我想要的文件）。因此，我的一些数据集也是使用pq.write_to_数据集编写的（有时我希望根据partition_filename_cb函数将数据重新分区到单个拼花地板文件中）。没有basename_模板，因此在编写数据集时，我的工作是使用元数据_收集器：files.append（fs.sep.join（[root_path，piece.row_group（0）.column（0）.file_path]）。这似乎相当方便，并且不需要我列出blob。我正在使用你的方法来处理ds.write_dataset（）数据集-希望有一天我能有一个标准的方法。酷-这对我来说很好，非常感谢-我不认为没有你我会走上这条路。因为我正在传递文件路径，所以我可以用它们生成一个数据集，然后使用fragment.partition_表达式生成一个过滤器表达式列表。。。做了很多工作，但它做到了我想要的