pyarrow - identify the fragments written or filters used when writing a parquet dataset?

user_5605073 0
Viewed 380

My use case is that I want to pass the file paths or filters to a task in Airflow as an xcom so that my next task can read the data which was just processed.

Task A writes a table to a partitioned dataset and a number of Parquet file fragments are generated --> Task B reads those fragments later as a dataset. I need to only read relevant data though, not the entire dataset which could have many millions of rows.

I have tested two approaches:

  1. List modified files right after I finish writing to the dataset. This will provide me with a list of paths which I can call ds.dataset(paths) on during my next task. I can use partitioning.parse() on these paths or check the fragments to get a list of filters used (frag.partition_expression)

A flaw with this is that I can have files being written in parallel to the same dataset.

  1. I can generate the filters used when writing the dataset by turning the table into a pandas dataframe, doing a groupby, and then constructing filters. I am not sure if there is a simpler approach to this. I can then use pq._filters_to_expression() on the results to create a usable filter.

This is not ideal since I need to fix certain data types which do not get saved properly as an Airflow xcom (no pickling so everything has to be in json format). Also, if I want to partition on a dictionary column, I might need to tweak this function.

    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

Any suggestions on which approach I should take, or if there is a better way to achieve my goal?

dask pyarrow
2 Answers

You can watch this issue (https://issues.apache.org/jira/browse/ARROW-10440) which does what you want I believe. In the meantime, you could use basename_template as a workaround.

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'))

This is just a rough sketch. You'd probably also want code to run at startup to see what the next free value of counter is. Or you could use a uuid instead of a counter.

user_2026940

A suggestion (not sure if this is optimal for your use case or not):

The key problem is the need to correctly select subset of the data, this can be 'fixed' upstream. The function/script that updates the big dataframe can contain a condition to save a temporary copy of data that is modified and satisfies some requirements in a separate (temporary) path. Then this file would be passed to the downstream tasks, which can delete the temporary file once it's processed.

user_106935960
Related