pyarrow Table Filtering -- huggingface

user_10839859 0
Viewed 453

I’m trying to filter a dataset based on the ids in a list. This approach is too slow. The dataset is an Arrow dataset. Import data from huggingface.

import numpy as np
from datasets import load_dataset, DatasetDict
from collections import Counter
import pyarrow as pa
import pandas as pd


responses = load_dataset('peixian/rtGender', 'responses', split = 'train')
# post_id_test_list contains list of ids
responses_test = responses.filter(lambda x: x['post_id'] in post_id_test_list)
pyarrow python
2 Answers

The dataset you get from load_dataset isn't an arrow Dataset but a hugging face Dataset. It is backed by an arrow table though.

Applying a lambda filter is going to be slow, if you want a faster vertorized operation you could try to modify the underlying arrow Table directly:

import pyarrow as pa
import pyarrow.compute as compute


table = responses.data

flags = compute.is_in(table['post_id'], value_set=pa.array(post_id_test_list, pa.int32()))
filtered_table = table.filter(flags)

filtered_respoonse = datasets.DataSet(filtered_table, response.info, response.split)

Though I'm not 100% sure if the last line is the correct way to recreate your dataset using an arrow table.

user_1095250

That almost got me there. As noted, the last line doesn't work, but I can convert to pandas / save out etc. Thanks!

import pyarrow as pa
import pyarrow.compute as compute


table = responses.data

flags = compute.is_in(table['post_id'], value_set=pa.array(post_id_test_list, pa.int32()))
filtered_table = table.filter(flags)
filtered_table.to_pandas()
user_108398590
Related