web-dev-qa-db-ja.com

pyarrowを使用してパーティション化された寄木細工のデータセットから特定のパーティションを読み取る

寄木細工の形式のやや大きい（〜20 GB）パーティション化されたデータセットがあります。 pyarrowを使用してデータセットから特定のパーティションを読み取りたいと思います。 pyarrow.parquet.ParquetDatasetでこれを達成できると思いましたが、そうではないようです。これは私が欲しいものを説明するための小さな例です。

ランダムなデータセットを作成するには：

from collections import OrderedDict
from itertools import product, chain
from uuid import uuid4
import os
from glob import glob

import numpy as np
import pandas as pd
import pyarrow as pa
from pyarrow.parquet import ParquetWriter, ParquetDataset


def get_partitions(basepath, partitions):
    """Generate directory hierarchy for a paritioned dataset

    data
    ├── part1=foo
    │   └── part2=True
    ├── part1=foo
    │   └── part2=False
    ├── part1=bar
    │   └── part2=True
    └── part1=bar
        └── part2=False

    """
    path_tmpl = '/'.join(['{}={}'] * len(partitions))  # part=value
    path_tmpl = '{}/{}'.format(basepath, path_tmpl)    # part1=val/part2=val

    parts = [product([part], vals) for part, vals in partitions.items()]
    parts = [i for i in product(*parts)]
    return [path_tmpl.format(*Tuple(chain.from_iterable(i))) for i in parts]


partitions = OrderedDict(part1=['foo', 'bar'], part2=[True, False])
parts = get_partitions('data', partitions)
for part in parts:
    # 3 columns, 5 rows
    data = [pa.array(np.random.Rand(5)) for i in range(3)]
    table = pa.Table.from_arrays(data, ['a', 'b', 'c'])
    os.makedirs(part, exist_ok=True)
    out = ParquetWriter('{}/{}.parquet'.format(part, uuid4()),
                        table.schema, flavor='spark')
    out.write_table(table)
    out.close()

パーティション1のすべての値を読み取り、パーティション2の場合はTrueのみを読み取りたいのですが、pandas.read_parquetを使用すると、それは不可能であり、常に列全体を読み取る必要があります。 pyarrowで次のことを試しました。

parts2 = OrderedDict(part1=['foo', 'bar'], part2=[True])
parts2 = get_partitions('data', parts2)
files = [glob('{}/*'.format(dirpath)) for dirpath in parts2]
files = [i for i in chain.from_iterable(files)]
df2 = ParquetDataset(files).read().to_pandas()

それも機能しません：

>>> df2.columns
Index(['a', 'b', 'c'], dtype='object')

これはpysparkで次のように簡単に実行できます。

def get_spark_session_ctx(appName):
    """Get or create a Spark Session, and the underlying Context."""
    from pyspark.sql import SparkSession
    spark = SparkSession.builder.appName(appName).getOrCreate()
    sc = spark.sparkContext
    return (spark, sc)


spark, sc = get_spark_session_ctx('test')
spark_df = spark.read.option('basePath', 'data').parquet(*parts2)
df3 = spark_df.toPandas()

以下に示すように：

>>> df3.columns
Index(['a', 'b', 'c', 'part1', 'part2'], dtype='object')

これはpyarrowまたはpandasで実行できますか、それともカスタム実装が必要ですか？

更新：Wesの要求に応じて、これは現在 [〜＃〜] jira [〜＃〜] にあります。

pythonparquetpyarrowapache-arrow
8
Quarky

質問：pyarrowを使用してパーティション化された寄木細工のデータセットから特定のパーティションを読み取るにはどうすればよいですか？

回答：今はできません。

https://issues.Apache.org/jira でこの機能をリクエストするApache Arrow JIRAを作成できますか？

これはpyarrowAPIでサポートできるはずですが、誰かがそれを実装する必要があります。ありがとうございました

5
Wes McKinney

Pyarrowバージョン0.10.0以降、filterskwargを使用してクエリを実行できます。あなたの場合、それは次のようになります：

import pyarrow.parquet as pq
dataset = pq.ParquetDataset('path-to-your-dataset', filters=[('part2', '=', 'True'),])
table = dataset.read()

参照

10
XiUpsilon