Parquet-bestanden schrijven met Haskell

Een voorbeeld:

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

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

Als je meer gedetailleerde controle over het Parquet-bestand nodig hebt, kun je writeParquetWithOptions gebruiken. Lees verder om te zien wat deze opties zijn en hoe ze het uiteindelijke bestand beïnvloeden.

Parquet voor Haskell

Om Haskell te laten samenwerken met het data-ecosysteem, moet het in staat zijn om de standaardformaten te begrijpen die in dat ecosysteem worden gebruikt. Lange tijd waren onze opties voor het serialiseren van data¹ in Haskell beperkt tot CSV, JSON, of het simpelweg dumpen van data in een ByteString. Hoewel CSV en JSON hun nut hebben, brengen ze aanzienlijke nadelen met zich mee naarmate het datavolume toeneemt. Ze zijn traag en omslachtig om te lezen, schrijven en bevragen, en elke eigen ontwikkelde indeling mist meestal de interoperabiliteit met standaard data science-tools.

Het Parquet-formaat offert eenvoud op voor efficiënte opslag en het bevragen van de opgeslagen data. We willen dat onze data hoge compressieratio's heeft en we willen het lezen van data die irrelevant is voor een specifieke query minimaliseren. De mogelijkheid om Parquet-bestanden te lezen en schrijven voor langetermijnopslag, verzending over het netwerk of voor interoperabiliteit met andere programma's — zeker gezien de universaliteit van Parquet in het data science-ecosysteem — is een zeer nuttig hulpmiddel dat we graag in een Dataframe-bibliotheek willen hebben.

Wat is Parquet?

(Sla deze sectie over als je de structuur van een Parquet-bestand al kent.)

Parquet-bestanden bestaan uit een reeks row groups, gevolgd door metadata aan het einde van het bestand met informatie die lezers in staat stelt om relevante column chunks en pagina's te lokaliseren. Het bevat ook nuttige informatie zoals statistieken en bloom-filters, zodat bijvoorbeeld een reader of query-planner kan beslissen of een specifieke row group wel of niet gelezen moet worden.

Elke row group is een verzameling van column chunks, waarvan elk hetzelfde aantal rijen bevat. Elke column chunk is een reeks datpagina's. Omdat elke column chunk een reeks pagina's is, en elke row group een reeks column chunks, ziet het uiteindelijke bestand er simpelweg uit als een reeks pagina's van elke kolom, satu achter de ander. We kunnen hier zinvol mee omgaan door de metadata te gebruiken om de offset en grootte van elke row group en column chunk te identificeren.

Een datpagina is de plek waar we daadwerkelijk alle data opslaan. Deze bestaat eerst uit de paginametadata, die onder andere de codering, het aantal waarden en de statistieken beschrijft. Er zijn feitelijk twee versies van de datpagina met subtiele verschillen.

Vervolgens hebben we definition levels en repetition levels. Dit is een goedkope codering van nullability (of een waarde nul kan zijn) en geneste structuren. Een gedetailleerde beschrijving hiervan valt buiten de scope van dit artikel; verwijs hiervoor naar het Dremel-paper. Voor onze doeleinden ondersteunen we momenteel alleen het schrijven van definition levels tot één, om nullable waarden aan te duiden.

Ten slotte hebben we onze feitelijk gecodeerde en gecomprimeerde data (afhankelijk van welke datpagina we gebruiken, comprimeren we alleen de data of zowel de definition/repetition levels als de data). De codering wordt per pagina bepaald, terwijl compressie per column chunk wordt bepaald.

De Implementatie

Hier volgt een technische bespreking van het ontwerp van de Parquet-writer en de gemaakte afwegingen.

Schrijfopties

Uit de beschrijving van het Parquet-formaat volgt dat een writer de juiste set aan parameters (de "fiddle factors") moet blootstellen, zodat de gebruiker de writer kan afstemmen op een Parquet-bestand dat efficiënt is voor de specifieke data. We willen een Parquet-writer die goed comprimeert met passend gekozen row groups en pagina's, zodat een reader effectief parallellisme, projection pushdown (alleen relevante column chunks lezen), predicate pushdown (alleen relevante row groups lezen / irrelevante pagina's negeren) en IO pushdown (alleen relevante bestanden lezen, uitgaande van data verdeeld over meerdere bestanden) kan toepassen.²

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

Zowel pageSize als rowGroupSize zijn de doelmatige grootte in bytes van elke pagina en elke row group. Maar elke column chunk in een row group moet hetzelfde aantal rijen bevatten. Afhankelijk van de specifieke data, de codering en het compressie-algoritme zal elke column chunk een ander aantal rijen bevatten voordat de doelgrootte wordt bereikt. Hetzelfde geldt voor elke pagina. Hoe zorgen we dan voor de doelgrootte van de pagina en row group, terwijl we hetzelfde aantal rijen in elke column chunk behouden?

We beschouwen zowel de pageSize als de rowGroupSize als een "best effort"; ze kunnen iets boven of onder de doelwaarde liggen. We verwerken onze kolommen in batches van batchRows en controleren de status van de row group na elke batch. Elke row group bevat dus een geheel veelvoud van batchRows rijen; pagina's zullen ook een geheel veelvoud van subBatchRows rijen bevatten (met uitzondering van de laatste pagina en de laatste column chunk). Sub-batching op paginaniveau stelt ons in staat om de hoeveelheid IORef boekhouding na elke schrijfactie te verminderen, wat aanzienlijke snelheidswinst oplevert.³

Geheugen

Vanwege de beperkingen beschreven in de vorige sectie is het moeilijk om vooraf te weten hoe groot onze buffers moeten zijn. Dit geldt des te meer voor column chunk-buffers: elke kolom zal een heel ander aantal rijen/pagina's in dezelfde hoeveelheid data passen, en we kunnen verwachten dat sommige buffers significant groter zullen zijn dan andere en de ruimte per row group zullen domineren.

We moeten ons geheugen dus kunnen uitbreiden. Een van de handigste manieren om ruw geheugen te beheren zonder terug te vallen op de FFI (Foreign Function Interface), is MutableByteArray.

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

Voor onze implementatie gebruiken we pinned ByteArrays, omdat we deze willen kunnen converteren naar een Ptr Word8 op het moment dat we ze moeten flushen naar een andere buffer (bijvoorbeeld een paginabuffer naar een column chunk-buffer) of naar een bestand (bij het flushen van een row group).

Het gebruik van pinned ByteArrays zorgt voor een kleine complicatie bij het uitbreiden van de geheugenbuffer. We mogen niet de grow-functie gebruiken die door Data.Primitive wordt geleverd; in plaats daarvan moeten we een nieuwe pinned ByteArray alloceren en de oude laten opruimen door de Garbage Collector (GC). Men zou zich zorgen kunnen maken over heap-fragmentatie omdat een enkel pinned object in een 4KB GHC-blok het hele blok in leven kan houden, maar we verwachten dat onze buffers veel groter zullen zijn dan 4KB. Bovendien zouden uitbreidingen zeldzaam moeten zijn, zeker na de eerste paar pagina's en de eerste RowGroup.

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 #-}

We hebben ook helperfuncties geschreven voor het schrijven van Word8s, Word32s, Word64s, Int32s, Int64s, Integers, Floats, Doubles en ByteStrings naar een buffer. Tot slot hebben we flushBufferToBuffer :: MemoryBuffer -> MemoryBuffer -> IO () en flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer -> IO ().⁴

De Kernlus

Hier volgt een high-level schets van de hoofdloop van de Parquet Writer. We laten aanzienlijke complexiteit weg om het ontwerp zo eenvoudig mogelijk uit te leggen. Voor de volledige implementatie kun je kijken in Writer.hs.

Op een hoog niveau kunnen we de Parquet-writer zien als een "effectful fold" over het dataframe⁵ (ook bekend als reduce of accumulate). Een fold kan worden gezien als een loop: we itereren over de rijen in ons dataframe. Folds zijn meestal puur, maar "effectful" betekent dat elke iteratie een zijeffect produceert; in ons geval het schrijven naar en vanuit het geheugen en naar de schijf. De fold produceert uiteindelijk één resultaat: de bestandsmetadata die we aan het einde van het bestand moeten toevoegen.

Eerst definiëren we de state die door de writer wordt geregen om dit te bereiken:

data ParquetWriterState = ParquetWriterState
{ outputFileHandle :: !WritableBinaryHandle -- newtype over een handle specifiek voor ons
, columnChunks :: !(VB.Vector ColumnChunkState) -- effectief onze row group buffer
, 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 -- Elke kolom heeft zijn eigen pageBuffer
, definitionLevels :: !DefLevels
, currentRowCount :: !(IORef Int)
}

-- In Encoder.hs
data Encoder = Encoder
{ encType :: !ThriftType
, convertedType :: !(Maybe ConvertedType)
, logicalType :: !(Maybe LogicalType)
, encodeValue :: !(MemoryBuffer -> Int -> Int -> IO (Int, Bool))
, finishValues :: !(MemoryBuffer -> Int -> IO Int) -- voor cleanup of post-processing
}

De buffers voor elke ColumnChunk worden geflusht naar het uiteindelijke bestand, terwijl de individuele paginabuffers worden geflusht naar ColumnChunk-buffers (of preciezer: ze worden gebruikt om de data te assembleren die uiteindelijk naar de ColumnChunk-buffer wordt geflusht, aangezien we de data samen met de paginametadata en definition levels moeten flushen).

Als we de ceremonie van het instellen van de initiële writerState en andere variabelen overslaan, is onze kernlus vrij eenvoudig:

loop :: Int -> IO ()
loop rowNum
  | rowNum >= endRow = pure ()
  | otherwise = do
      let batchEnd = rowNum + min interval (endRow - rowNum)
      writeBatch rowNum batchEnd
      size <- bufferedSize writerState.columnChunks
      if (size >= options.rowGroupSize)
        then do
          flushRowGroup options writerState -- Schrijf naar bestand
          loop batchEnd
        else loop batchEnd

En voor 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

De scratchBuffer is een herbruikbare buffer die we later gebruiken om de pagina te assembleren; de individuele paginabuffers bevatten namelijk alleen de waarden, terwijl onze definition levels in een andere buffer leven en de paginametadata nog niet bestaat. Het is de scratchBuffer die uiteindelijk naar een ColumnChunk-buffer wordt geflusht.

De interne werking van writeRows lijkt enigszins op de loop die we eerder zagen:

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.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 $
        pushDef page.definitionLevels (if notNull then 1 else 0)
      rowWriterLoop options columnChunkState end size position' (row + 1)

Let op dat we voor onze writer alleen definition levels tot 1 afhandelen. We ondersteunen geen definition levels groter dan 1 en repetition levels groter dan 0, omdat we momenteel specifiek alleen dataframes met platte schema's willen ondersteunen.

Wanneer een pagina vol is, kunnen we deze flushen naar de parent columnChunkState.buffer:

pageResidency <- readIORef page.pageBuffer.positionRef
defResidency <- readIORef page.definitionLevels.dlBuf.positionRef
when (pageResidency + defResidency >= options.pageSize)
  (flushPage options scratchBuffer columnChunkState)

Toekomstig werk

Momenteel ondersteunt de Parquet Writer alleen:

• Snappy en Uncompressed compressie.
• Plain Encoding.
• Definition levels tot 1.

De volgende stap is het ondersteunen van het volledige scala aan compressies en coderingen die door Parquet worden aangeboden. De huidige Writer is ook volledig single-threaded; we kunnen onderzoeken om concurrency en/of parallellisme toe te voegen, mits daar een prestatie-rechtvaardiging voor is.

Momenteel worden optionele metadata-velden, zoals statistieken voor kolommen/pagina's en bloom-filters, niet geregistreerd. Zoals eerder vermeld, zijn deze zeer nuttig voor Parquet-readers om het bevragen van de data te optimaliseren.

Tot slot leven momenteel al onze buffers in het geheugen terwijl ze wachten om naar de schijf te worden geflusht. In systemen met beperkt geheugen kunnen we bij het schrijven van Parquet-bestanden met zeer grote pagina's of row groups uit het geheugen raken (een volledige row group moet namelijk volledig in het geheugen worden gehouden). We plannen daarom een "two-pass" strategie waarbij we veel kleinere in-memory buffers gebruiken en naar tijdelijke bestanden op schijf schrijven in plaats van alles in het geheugen te houden.

***

Voetnoten:

¹ Wanneer we in dit artikel over "data" spreken, bedoelen we meestal grote hoeveelheden kolomgebaseerde data die men gewoonlijk in een dataframe plaatst.

² De bovenstaande optimalisaties vereisen ook de implementatie van meerdere coderingen, statistieken, bloom-filters, enz. De huidige implementatie richt zich op het operationeel krijgen van de basis-writer; de optionele delen van Parquet-bestanden zijn gepland als toekomstig werk.

³ Er is een mogelijkheid dat, indien kolommen individuele entries hebben die disproportioneel groot zijn, onze sub-batching/batching de doel pageSize en/of rowGroupSize met orders van grootte overschrijdt. Een mechanisme om dit effect te beperken moet mogelijk worden geïmplementeerd als gebruikers hier tegenaan lopen.

⁴ De implementatie van flushBufferToFile leidde tot een interessant inzicht: standaard schrijven we bij het wegschrijven naar een bestand eigenlijk naar de OS-paginacheche. Als de ratio van dirty pages een bepaalde waarde overschrijdt, knijpt het OS het proces dat de dirty pages produceert af totdat de writeback is bijgehouden. Omdat we row groups synchroon schrijven, worden we gedwongen te wachten tot de Kernel alles heeft gesorteerd (hPutBuf gebruikt intern write(2)). Daarom is het raadzaam om data in grote genoeg chunks te schrijven zodat we meer data met minder syscalls schrijven. Na tests met dd is 256 KiB chunks als een redelijke standaard gekozen.

-- voor gebruik in een single-threaded writer
flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer -> IO ()
flushBufferToFile (WritableBinaryHandle h) buffer = do
  array <- readIORef buffer.arrayRef
  position <- readIORef buffer.positionRef
  withMutableByteArrayContents array $ \ptr -> do
    let chunkSize = 262144
    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

⁵ We spreken over een "effectful fold over het dataframe", maar als je naar de code kijkt, gebruiken we geen foldM. De code is grotendeels recursief en er wordt een forM gebruikt. Dit is de lokale oplossing om het batching- en sub-batching-proces handig te maken. De bewering dat het een effectful fold is, blijft echter accuraat, aangezien forM zelf een speciaal soort fold is en folds abstracties zijn over recursie (zie het paper van Graham Hutton: “A Tutorial on the Universality and Expressiveness of Fold”).