HaskellでParquetファイルを書き出す

Haskell で Parquet ファイルを書き込む

DataHaskell/Dataframe に Parquet ライターを実装しました。最も一般的なケースでは、非常に簡単に使用できます。データフレームを writeParquet 関数に渡すだけで、妥当なデフォルト設定で Parquet ファイルが書き込まれます。例:

import qualified DataFrame as D
import qualified DataFrame.Functions as F
import DataFrame (as, (|>))

$(D.declareColumnsFromParquetFile "sales_data.parquet")

main = do
    sales <- D.readParquet "sales_data.parquet"
    
    sales
        |> D.groupBy   ["product"]
        |> D.aggregate [ F.sum amount `as` "total" 
                       , F.countAll `as` "orders"
                       ]
        |> D.writeParquet "total_orders.parquet"

Parquet ファイルに対してより詳細な制御が必要な場合は、writeParquetWithOptions を使用することになります。オプションの内容と、それが最終的なファイルにどのように影響するかについては、読み進めてください。

Haskell がデータエコシステムと相互運用するためには、そのエコシステムで使用されている標準フォーマットを理解できる必要があります。長らく、Haskell でデータをシリアライズする1ための選択肢は、CSV、JSON、または単に ByteString にダンプするだけでした。CSV と JSON にはそれぞれの用途がありますが、データ量が増加するにつれて、かなりの欠点があります。読み込み、書き込み、クエリが遅く煩雑であり、独自のカスタムフォーマットは標準的なデータサイエンスツールとの相互運用性に欠ける傾向があります。

Parquet フォーマットは、単純さと引き換えに、そこに格納されたデータの効率的なストレージとクエリを実現します。データの圧縮率を高く保ち、特定のクエリに関連しないデータの読み込みを最小限に抑えたいと考えています。長期保存、ネットワーク送信、または他のプログラムとの相互運用（特にデータサイエンスエコシステムにおける Parquet の普遍性を考えると）のために Parquet ファイルを読み書きできる能力は、データフレームライブラリに欲しい非常に便利なツールです。

Parqué?

Parquet ファイルの構造をすでに知っている場合は、このセクションをスキップしても構いません。

Parquet ファイルは、一連の行グループと、ファイル末尾のメタデータで構成されており、メタデータには、関連する列チャンクとページを特定するための情報が含まれています。また、統計情報やブルームフィルターなどの有用な情報も含まれているため、例えば、リーダーやクエリプランナーは特定の行グループを読み取るべきかどうかを判断できます。

各行グループは列チャンクの集まりであり、各列チャンクは同じ数の行を含んでいます。各列チャンクはデータページのシリーズです。各列チャンクはページのシリーズであり、各行グループは列チャンクのシリーズであるため、最終的なファイルは、各列からのページのシリーズが次々と続くように見えます。メタデータを使用して、各行グループと列チャンクのオフセットとサイズを特定することで、すべてを理解することができます。

データページは、実際のデータを格納する場所です。まず、ページメタデータ（エンコーディング、値の数、統計情報などを記述したもの）で構成されます。データページには、実際には2つの バージョンがあり、微妙な違いがあります。

次に、定義レベルと繰り返しレベルがあります。これらは、null 値の許容度とネスト構造を効率的にエンコードしたものです。定義レベルと繰り返しレベルの詳細な説明は、この記事の範囲外です。詳細については、Dremel paperを参照してください。ここでは、null 値を示すために、定義レベルを 1 までのみサポートしています。

最後に、実際のエンコードおよび圧縮されたデータがあります（使用するデータページによっては、データのみ、または定義レベル/繰り返しレベルとデータの両方を圧縮します）。エンコーディングはページごとに決定され、圧縮は列チャンクごとに決定されます。

実装

以下は、Parquet ライターの設計と、私たちが下したトレードオフに関する技術的な議論です。

書き込みオプション

上記の Parquet ファイルの形式の説明から、ライターは、ユーザーが特定のデータに適用される Parquet ファイルを効率的に生成するようにライターを調整できるように、適切な調整係数（fiddle factors）のセットを公開する必要があることがわかります。つまり、適切なサイズの行グループとページでうまく圧縮される Parquet ライターを望んでおり、リーダーが並列処理、射影プッシュダウン（関連する列チャンクのみを読み取る）、述語プッシュダウン（関連する行グループのみを読み取る/無関係なページを削除する）、IO プッシュダウン（データが複数のファイルにチャンク化されていると仮定して、関連するファイルのみを読み取る）などを効果的に適用できるようにしたいと考えています2。

data ParquetWriteOptions = ParquetWriteOptions
    { pageSize :: !Int
    , rowGroupSize :: !Int
    , batchRows :: !Int
    , subBatchRows :: !Int
    , compressionCodec :: !CompressionCodec
    , strategy :: !WriterStrategy
    , maxRowsPerFile :: !(Maybe Int)
    }

pageSize と rowGroupSize は、各ページと各行グループのターゲットサイズ（バイト単位）です。しかし、行グループ内の各列チャンクは同じ数の行を含んでいる必要があり、エンコードされる特定のデータ、エンコーディング、使用する圧縮アルゴリズムによっては、各列チャンクはターゲットサイズに達する前に異なる数の行を保持します。各ページについても同様です。では、ページサイズのターゲットと行グループサイズのターゲットを確実にし、各列チャンクに同じ数の行を含めるにはどうすればよいでしょうか？

ターゲット pageSize と rowGroupSize の両方を最善努力（best effort）と見なす必要があります。ターゲットよりやや大きい場合も小さい場合もあります。バッチサイズ batchRows のバッチで列を処理し、各バッチ後に列グループの状態を確認します。したがって、各行グループには batchRows の整数倍の行が含まれます。ページも（最後のページと最後の列チャンクを除いて）subBatchRows の整数倍の行を含むことになります。ページレベルでのサブバッチ処理により、書き込みごとに必要となる IORef のブックキーピングを削減でき、大幅な速度向上が得られます3。

メモリ

前のセクションで説明した制約のため、バッファを事前にどれくらいの大きさにするかを把握するのは困難です。これは列チャンクバッファについても同様です。各列は、同じデータ量で非常に異なる数の行/ページを格納でき、行グループあたりのスペースを支配する列チャンクバッファが他のものよりかなり大きくなることが予想されます。

そのため、メモリを拡張できる必要があります。FFI に頼らずに生のメモリを扱う便利な方法の 1 つは MutableByteArray です。

data MemoryBuffer = MemoryBuffer
    { arrayRef :: !(IORef (MutableByteArray RealWorld))
    , positionRef :: !(IORef Int)
    }

実装では、ピン留めされた ByteArray を使用します。これは、ページバッファを列チャンクバッファにフラッシュする場合や、行グループをファイルにフラッシュする場合のように、Ptr Word8 に変換したい場合があるためです。

ピン留めされた ByteArray を使用すると、メモリバッファを拡張する際にわずかな複雑さが生じます。Data.Primitive が提供する grow 関数を使用することはできません。代わりに、新しいピン留めされた ByteArray を割り当て、古いものは GC に任せる必要があります。単一のピン留めされたオブジェクトが 4KB の GHC ブロックでブロック全体を生き続けさせる可能性があるため、ヒープフラグメンテーションを心配するかもしれませんが、バッファは 4KB よりはるかに大きくなることが予想されます。さらに、特に最初の数ページと最初の行グループの後では、拡張はまれであるべきです。

ensureCapacity :: MemoryBuffer -> Int -> IO (MutableByteArray RealWorld)
ensureCapacity buffer needed = do
    array <- readIORef buffer.arrayRef
    maxSize <- getSizeofMutableByteArray array
    if needed <= maxSize
        then pure array
        else do
            position <- readIORef buffer.positionRef
            grown <- newPinnedByteArray (needed + (needed `div` 2)) -- 新しいバッファを少し大きめに確保
            copyMutableByteArray grown 0 array 0 position -- 既存の内容をコピー
            writeIORef buffer.arrayRef grown
            pure grown
{-# INLINE ensureCapacity #-}

また、Word8、Word32、Word64、Int32、Int64、Integer、Float、Double、および ByteString をバッファに書き込むためのヘルパー関数も作成しました。最後に、flushBufferToBuffer :: MemoryBuffer -> MemoryBuffer -> IO () と flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer -> IO () があります4。

コアループ

ここでは、Parquet ライターのメインループのハイレベルなスケッチを提供します。設計をできるだけシンプルに説明するために、かなりの複雑さは省略しています。完全な実装については、Writer.hs を参照してください。

ハイレベルでは、Parquet ライターをデータフレームに対する効果的なフォールド（他の言語では reduce、よりまれには accumulate とも呼ばれる）と考えることができます。フォールドはループと考えることができ、データフレームの行を反復処理します。フォールドは通常純粋ですが、「効果的」とは、各反復が何らかの副作用（この場合はメモリへの読み書きとディスクへの書き込み）を生成することを意味します。最後に、フォールドは反復から最終的な単一の結果を生成します。これは、ファイル末尾に追加する必要があるファイルメタデータです。

まず、これを達成するためにライターにスレッドされる状態を定義する必要があります。

data ParquetWriterState = ParquetWriterState
    { outputFileHandle :: !WritableBinaryHandle -- 私たちが書き込むためのハンドル専用の newtype
    , columnChunks :: !(VB.Vector ColumnChunkState) -- 実質的に私たちの行グループバッファ
    , currentFileOffsetRef :: !(IORef Int64)
    , scratchBuffer :: !MemoryBuffer
    , rowGroupMetadataRef :: !(IORef [RowGroup])
    , rowNumberRef :: !(IORef Int)
    }

data ColumnChunkState = ColumnChunkState
    { columnName :: !T.Text
    , nullable :: !Bool
    , schema :: !SchemaElement
    , encoder :: !Encoder
    , buffer :: !MemoryBuffer
    , uncompressedBufferSize :: !(IORef Int64)
    , pageState :: !PageState
    }

data PageState = PageState
    { pageBuffer :: !MemoryBuffer -- 各列には独自の pageBuffer があります
    , definitionLevels :: !DefLevels
    , currentRowCount :: !(IORef Int)
    }
    
-- Encoder.hs 内

data Encoder = Encoder
    { encType :: !ThriftType
    , convertedType :: !(Maybe ConvertedType)
    , logicalType :: !(Maybe LogicalType)
    , encodeValue :: !(MemoryBuffer -> Int -> Int -> IO (Int, Bool)) -- 値をエンコードし、新しい位置と null でなかったかを示すブール値を返す
    , finishValues :: !(MemoryBuffer -> Int -> IO Int) -- クリーンアップや後処理が必要な場合
    }

各 ColumnChunk のバッファは最終ファイルにフラッシュされ、個々のページバッファは ColumnChunk バッファにフラッシュされます（より正確には、ページメタデータと定義レベルとともにデータをフラッシュする必要があるため、最終的に ColumnChunk バッファにフラッシュされるデータを組み立てるために使用されます）。

初期 writerState やその他の変数の設定の儀式を省略すると、コア ループは非常にシンプルになります。

loop :: Int -> IO ()
loop rowNum
    | rowNum >= endRow = pure ()
    | otherwise = do
        let batchEnd = rowNum + min interval (endRow - rowNum)
        writeBatch rowNum batchEnd
        size <- bufferedSize writerState.columnChunks
        when
            (size >= options.rowGroupSize)
            (flushRowGroup options writerState) -- ファイルに書き込む
        loop batchEnd

次に writeBatch です。

writeBatch :: Int -> Int -> IO ()
writeBatch rowNum batchEnd
    | rowNum >= batchEnd = pure ()
    | otherwise = do
        let count = min options.subBatchRows (batchEnd - rowNum)
        forM_ writerState.columnChunks (writeRows options scratchBuffer rowNum count)
        modifyIORef writerState.rowNumberRef (+ count)
        writeBatch (rowNum + count) batchEnd
        

scratchBuffer は再利用可能なバッファであり、後でページを組み立てるために使用されます。個々のページバッファは値のみを保持しますが、定義レベルは別のバッファにあり、ページメタデータはまだ存在しません。最終的に ColumnChunk バッファにフラッシュされるデータが scratchBuffer です。

writeRows の内部動作は、前述の loop と似ています。

rowWriterLoop !options !columnChunkState !end !size !position !row
    | row >= end = writeIORef columnChunkState.pageState.pageBuffer.positionRef position
    | position + options.pageSize > size = do -- ページがいっぱいになりそうなら、新しいページを確保
        let page = columnChunkState.pageState
        writeIORef page.pageState.pageBuffer.positionRef position
        arr' <- ensureCapacity
                    page.pageBuffer
                    (position + max
                                    options.pageSize
                                    ((end - row) * 64) -- 次の行のために十分なスペースを確保
                    )
        size' <- getSizeofMutableByteArray arr' -- 新しいバッファのサイズを取得
        rowWriterLoop options columnChunkState end size' position row
    | otherwise = do
        let page = columnChunkState.pageState
            encode = columnChunkState.encoder.encodeValue
        (position', notNull) <- encode page.pageBuffer position row -- 値をエンコード
        when columnChunkState.nullable $ -- null 許容の場合
            pushDef page.definitionLevels (if notNull then 1 else 0) -- 定義レベルをプッシュ
        rowWriterLoop options columnChunkState end size position' (row + 1)

なお、ライターでは定義レベルを 1 までのみ扱い、定義レベルが 1 より大きい場合や繰り返しレベルが 0 より大きい場合はサポートしていません。これは、現在、フラットなスキーマを持つデータフレームのみをサポートすることを特に目的としているためです。

ページがいっぱいになると、親の columnChunkState.buffer にフラッシュできます。

pageResidency <- readIORef page.pageBuffer.positionRef
defResidency <- readIORef page.definitionLevels.dlBuf.positionRef
when
    (pageResidency + defResidency >= options.pageSize) -- ページサイズに達したら
    (flushPage options scratchBuffer columnChunkState) -- ページをフラッシュ

今後の作業

現在、Parquet ライターは以下のみをサポートしています。

• Snappy および圧縮なしの圧縮
• Plain エンコーディング
• 定義レベルは 1 までのみ

次のステップとしては、まず Parquet が提供するすべての圧縮およびエンコーディングをサポートすることです。現在のライターは完全にシングルスレッドであるため、限定的な方法で並列処理や並列性を追加することを検討するかもしれません。ただし、パフォーマンス上の正当な理由がある場合に限ります。

現在、列/ページごとの統計情報やブルームフィルターなどのオプションのメタデータフィールドは記録されていません。前述したように、これらは Parquet リーダーが Parquet ファイル内のデータをクエリする際に非常に役立ちます。

最後に、現在、すべてのバッファはディスクにフラッシュされるまでメモリ内にあります。特に大きなページ/行グループを必要とする Parquet ファイルをメモリ制約のあるシステムで書き込む場合、メモリ不足になる可能性があります（行グループ全体をメモリ全体に保持する必要があります）。そのため、より小さなインメモリバッファを使用し、データをメモリ内に保持する代わりにディスク上のテンポラリファイルに書き込む 2 パス戦略を実装する予定です。

---

1. この記事でデータについて話すとき、通常はデータフレームに入れるような大量の列指向データを意味します。↩

2. 上記の最適化には、複数のエンコーディング、統計情報、ブルームフィルターなどの実装も必要です。現在の実装は、基本的なライターを稼働させることに焦点を当てており、Parquet ファイルのオプション部分は将来の作業として計画されています。↩

3. 個々のエントリが不釣り合いに大きい列がある場合、サブバッチ処理/バッチ処理がターゲット pageSize および/または rowGroupSize を桁違いに超える可能性があります。ユーザーがこの問題に遭遇し始めた場合、この効果を軽減するためのメカニズムを実装する必要があるかもしれません。↩

4. flushBufferToFile の実装は、やや興味深い穴に私を導きました。私の乏しい理解によれば、ファイルへの書き込み時に実際には OS のページキャッシュに書き込んでおり、ページをダーティとしてマークし、ダーティページの比率が特定の値を上回ると、OS は書き戻しが追いつくのを待っている間にプロセスをスロットルします。行グループを同期的に書き込んでいるため、カーネルが物事を整理するのを待つことを余儀なくされます（hPutBuf は内部的に write(2) を使用します。もし私が間違っていなければ）。

したがって、システムコールを少なくしてより多くのデータを書き込むために、十分な大きさのチャンクでデータを書き込むことが重要です（ただし、非常に大きなチャンクはリターンの逓減があります）。私のコンピューターで dd を使用してテストを実行した後、Parquet Writer の妥当なデフォルトとして 256 KiB チャンクに落ち着きました。

-- シングルスレッドライターでの使用のため
flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer -> IO ()
flushBufferToFile (WritableBinaryHandle h) buffer = do
    array <- readIORef buffer.arrayRef
    position <- readIORef buffer.positionRef
    withMutableByteArrayContents array $ \ptr -> do
        let chunkSize = 262144 -- 256 KiB
            go offset
                | offset >= position = pure ()
                | otherwise = do
                    let n = min chunkSize (position - offset)
                    -- ハンドルはブロックバッファリングされていません
                    hPutBuf h (ptr `plusPtr` offset) n
                    go (offset + n)
        go 0
    writeIORef buffer.positionRef 0

ディスクへの直接書き込み（カーネルの仲介なし）を、クロスプラットフォーム対応の追加の労力を払ってでも望む場合、O_DIRECT が beckons します。マルチスレッドとクラッシュによるデータ損失の回避を考慮すると、穴はさらに深くなります。↩

5. 「データフレームに対する効果的なフォールド」と言っていますが、コードを見ると実際には foldM を使用していません。ほとんどのコードは再帰的であり、forM_ が使用されています。これは、ライターに必要なバッチ処理とサブバッチ処理を容易にするために着地したローカルミニマムです。しかし、効果的なフォールドであるという記述は、forM_ 自体が一種のフォールドであり、フォールドは再帰を抽象化する（Graham Hutton の論文「A Tutorial on the Universality and Expressiveness of Fold」を参照）ため、依然として正確です。↩