wireform-parquet-0.1.0.0: bench/SelectionThroughput.hs

-- Shared PyArrow fixtures and independent expected values are generated by
-- scripts/parquet_selection_fixtures.py; see the benchmark manifest.
module Main (main) where

import Arrow.Column qualified as AC
import Arrow.Column.Buffer qualified as Buffer
import Arrow.Types qualified as AT
import Control.DeepSeq (NFData, force)
import Control.Exception (evaluate)
import Control.Monad (unless)
import Criterion.Main (bench, bgroup, defaultMain, nf)
import Data.ByteString qualified as BS
import Data.Int (Int32, Int64)
import Data.Maybe (isJust)
import Data.Text qualified as Text
import Data.Vector qualified as V
import Data.Vector.Primitive qualified as VP
import Data.Vector.Storable qualified as VS
import GHC.Generics (Generic)
import Parquet.Arrow qualified as Arrow
import Parquet.HighLevel qualified as HighLevel
import Parquet.PageIndex qualified as PageIndex
import Parquet.Predicate qualified as Predicate
import Parquet.Read qualified as Read
import Parquet.Types qualified as P
import System.Environment (getEnv, lookupEnv)
import System.FilePath ((</>))
import Text.Read (readMaybe)


data Fixture = Fixture
  { fixtureName :: !String
  , fixtureBytes :: !BS.ByteString
  , fixturePredicate :: !Predicate.PColPredicate
  , fixtureFile :: !Read.ParquetFile
  , fixtureField :: !AT.Field
  , fixtureCodec :: !P.Compression
  , fixtureSelection :: !Read.PageSelection
  }


data RawValues = RawInt64 !(VP.Vector Int64) | RawBytes !(V.Vector BS.ByteString)
  deriving stock (Eq, Generic)
  deriving anyclass (NFData)


data RawColumn = RawColumn
  { rawRepetition :: !(VP.Vector Int32)
  , rawDefinition :: !(VP.Vector Int32)
  , rawValues :: !RawValues
  }
  deriving stock (Eq, Generic)
  deriving anyclass (NFData)


data BufferColumn = BufferColumn
  { bufferRepetition :: !(VP.Vector Int32)
  , bufferDefinition :: !(VP.Vector Int32)
  , bufferOffsets :: !(VS.Vector Int32)
  , bufferValues :: !BS.ByteString
  }
  deriving stock (Eq, Generic)
  deriving anyclass (NFData)


-- Time the native buffer readers used by Arrow, with metadata/index planning
-- prepared. The legacy vector group remains an independent control.
selectedBuffers :: Fixture -> BS.ByteString -> BufferColumn
selectedBuffers fixture bytes = either error id $ case AT.fieldType field of
  AT.AInt 64 True
    | nullable -> leveled scalar (Read.readScalarBuffersSelectedLevels P.PTInt64 codec 0 1 bytes selection)
    | otherwise -> dense scalar (Read.readScalarBuffersSelectedValues P.PTInt64 codec bytes selection)
  AT.AUtf8
    | nullable -> leveled binary (Read.readBinaryBuffersSelectedLevels codec 0 1 bytes selection)
    | otherwise -> dense binary (Read.readBinaryBuffersSelectedValues codec bytes selection)
  _ -> Left "unsupported buffer selection benchmark type"
  where
    field = fixtureField fixture
    nullable = AT.fieldNullable field
    codec = fixtureCodec fixture
    selection = fixtureSelection fixture
    scalar reps defs values = BufferColumn reps defs VS.empty (Read.scalarData values)
    binary reps defs values = BufferColumn reps defs (Read.binaryOffsets values) (Read.binaryData values)
    dense wrap decoded = wrap VP.empty VP.empty <$> decoded
    leveled wrap decoded = do
      levels <- decoded
      Right (wrap (Read.clRepetition levels) (Read.clDefinition levels) (Read.clValues levels))


expectedBuffers :: RawColumn -> BufferColumn
expectedBuffers raw = case rawValues raw of
  RawInt64 values -> BufferColumn reps defs VS.empty (Buffer.storableToBytes (VP.convert values :: VS.Vector Int64))
  RawBytes values -> case AC.fromByteStrings values of
    AC.ColBinary Nothing offsets bytes -> BufferColumn reps defs offsets bytes
    _ -> error "expected dense binary reference buffers"
  where
    reps = rawRepetition raw
    defs = rawDefinition raw


-- Profile the legacy public vector readers, independently of the contiguous
-- binary-buffer path used by Arrow. Only metadata and selection planning are
-- prepared: headers, decompression, levels, dictionaries and values are
-- decoded on every iteration, without constructing an Arrow column.
selectedRaw :: Fixture -> BS.ByteString -> RawColumn
selectedRaw fixture bytes = either error id $
  case AT.fieldType field of
    AT.AInt 64 True
      | nullable -> leveled RawInt64 (Read.readInt64SelectedLevels codec 0 1 bytes selection)
      | otherwise -> dense RawInt64 (Read.readInt64SelectedValues codec bytes selection)
    AT.AUtf8
      | nullable -> leveled RawBytes (Read.readByteArraySelectedLevels codec 0 1 bytes selection)
      | otherwise -> dense RawBytes (Read.readByteArraySelectedValues codec bytes selection)
    _ -> Left "unsupported raw selection benchmark type"
  where
    field = fixtureField fixture
    nullable = AT.fieldNullable field
    codec = fixtureCodec fixture
    selection = fixtureSelection fixture
    dense wrap decoded = RawColumn VP.empty VP.empty . wrap <$> decoded
    leveled wrap decoded = do
      levels <- decoded
      Right (RawColumn (Read.clRepetition levels) (Read.clDefinition levels) (wrap (Read.clValues levels)))


expectedRaw :: AT.Field -> AC.ColumnArray -> Either String RawColumn
expectedRaw field column = case AT.fieldType field of
  AT.AInt 64 True -> do
    array <- maybe (Left "expected INT64 array") Right (AC.asPrim AC.PInt64 column)
    Right (fromRows (RawInt64 . VP.convert) (V.generate count (AC.unsafePrimAt array)))
  AT.AUtf8 -> do
    array <- maybe (Left "expected byte array") Right (AC.asBinary column)
    Right (fromRows RawBytes (V.generate count (AC.unsafeBytesAt array)))
  _ -> Left "unsupported expected raw benchmark type"
  where
    count = AC.columnLength column
    fromRows wrap rows =
      RawColumn
        (if AT.fieldNullable field then VP.replicate count 0 else VP.empty)
        (if AT.fieldNullable field then VP.generate count (\i -> if isJust (rows V.! i) then 1 else 0) else VP.empty)
        (wrap (V.mapMaybe id rows))


data SelectionPlan = SelectionPlan
  { planCodec :: !P.Compression
  , planSelection :: !Read.PageSelection
  }


prepareSelection :: Predicate.PColPredicate -> Read.ParquetFile -> Either String SelectionPlan
prepareSelection predicate file = do
  bounds <- PageIndex.readColumnIndex file 0 0 >>= maybe (Left "missing column index") Right
  offsets <- PageIndex.readOffsetIndex file 0 0 >>= maybe (Left "missing offset index") Right
  let footer = Read.pfFooter file
  group <- maybe (Left "missing row group") Right (P.fmRowGroups footer V.!? 0)
  metadata <- maybe (Left "missing column metadata") Right (P.rgColumns group V.!? 0 >>= P.ccMetadata)
  context <- maybe (Left "missing statistics context") Right (Predicate.statisticsContext footer 0)
  let keep = (== Predicate.PMaybeKeep) <$> Predicate.evalPagesByColumnIndex context bounds predicate
  Right (SelectionPlan (P.cmCodec metadata) (Read.PageSelection (P.oiPageLocations offsets) keep (P.cmDictionaryPageOffset metadata) (Just (P.rgNumRows group))))


selected :: Predicate.PColPredicate -> BS.ByteString -> AC.ColumnArray
selected predicate bytes = either error id $ do
  file <- HighLevel.decodeParquet HighLevel.defaultReadOptions bytes
  schema <- Arrow.parquetFileArrowSchemaChecked file
  field <- maybe (Left "selection fixture has no column") Right (AT.arrowFields schema V.!? 0)
  snd <$> Arrow.readParquetColumnWithPagePruning file 0 0 field predicate


-- File/schema planning is forced and value-checked during fixture loading.
-- This still times index parsing, page selection, decoding and Arrow assembly.
selectedPlanned :: Predicate.PColPredicate -> AT.Field -> Read.ParquetFile -> AC.ColumnArray
selectedPlanned predicate field file =
  either error snd $
    Arrow.readParquetColumnWithPagePruning file 0 0 field predicate


-- Force both index readers and every pruning decision without decoding pages.
-- Summing retained page row ranges gives a check against the independent output.
selectedIndexRows :: Predicate.PColPredicate -> Read.ParquetFile -> Int64
selectedIndexRows predicate file = either error id $ do
  bounds <- PageIndex.readColumnIndex file 0 0 >>= maybe (Left "missing column index") Right
  offsets <- PageIndex.readOffsetIndex file 0 0 >>= maybe (Left "missing offset index") Right
  let metadata = Read.pfFooter file
      locations = P.oiPageLocations offsets
  group <- maybe (Left "missing row group") Right (P.fmRowGroups metadata V.!? 0)
  context <- maybe (Left "missing statistics context") Right (Predicate.statisticsContext metadata 0)
  let keep = (== Predicate.PMaybeKeep) <$> Predicate.evalPagesByColumnIndex context bounds predicate
      rows total i retained
        | not retained = total
        | otherwise =
            let start = P.plFirstRowIndex (locations V.! i)
                end = maybe (P.rgNumRows group) P.plFirstRowIndex (locations V.!? (i + 1))
            in total + end - start
  if V.length locations /= V.length keep
    then Left "page index lengths disagree"
    else Right (V.ifoldl' rows 0 keep)


loadFixtures :: IO [Fixture]
loadFixtures = do
  root <- getEnv "WIREFORM_PARQUET_SELECTION_FIXTURES"
  plans <- lines <$> readFile (root </> "selection-plan.txt")
  traverse (load root . words) plans
  where
    load root [name, typ, literal] = do
      value <- case typ of
        "text" -> pure (Predicate.PVText (Text.pack literal))
        "int" -> maybe (fail "invalid integer selection pivot") (pure . Predicate.PVInt64) (readMaybe literal)
        _ -> fail "invalid selection fixture type"
      bytes <- BS.readFile (root </> (name ++ ".parquet"))
      expectedBytes <- BS.readFile (root </> (name ++ "-expected.parquet"))
      expected <- evaluate $ force $ either error id $ do
        file <- HighLevel.decodeParquet HighLevel.defaultReadOptions expectedBytes
        schema <- Arrow.parquetFileArrowSchemaChecked file
        field <- maybe (Left "expected fixture has no column") Right (AT.arrowFields schema V.!? 0)
        Arrow.readParquetColumn file 0 0 field
      let predicate = Predicate.PGtEq value
      actual <- evaluate (force (selected predicate bytes))
      unless (actual == expected) (fail ("selection value mismatch: " ++ name))
      file <- either fail pure (HighLevel.decodeParquet HighLevel.defaultReadOptions bytes)
      schema <- either fail pure (Arrow.parquetFileArrowSchemaChecked file)
      field <- maybe (fail "selection fixture has no column") pure (AT.arrowFields schema V.!? 0)
      planned <- evaluate (force (selectedPlanned predicate field file))
      unless (planned == expected) (fail ("planned selection value mismatch: " ++ name))
      indexedRows <- evaluate (selectedIndexRows predicate file)
      unless (indexedRows == fromIntegral (AC.columnLength expected)) (fail ("index row count mismatch: " ++ name))
      plan <- either fail pure (prepareSelection predicate file)
      let fixture = Fixture name bytes predicate file field (planCodec plan) (planSelection plan)
      raw <- evaluate (force (selectedRaw fixture bytes))
      expectedValues <- either fail (evaluate . force) (expectedRaw field expected)
      unless (raw == expectedValues) (fail ("raw selection value/level mismatch: " ++ name))
      buffers <- evaluate (force (selectedBuffers fixture bytes))
      referenceBuffers <- evaluate (force (expectedBuffers expectedValues))
      unless (buffers == referenceBuffers) (fail ("native buffer selection value/level mismatch: " ++ name))
      pure fixture
    load _ _ = fail "invalid selection plan"


main :: IO ()
main = do
  fixtures <- loadFixtures
  profile <- (== Just "1") <$> lookupEnv "WIREFORM_PARQUET_PROFILE_PHASES"
  let planned =
        if profile
          then
            [ bgroup "selected pages planned" (plannedBenchmark <$> fixtures)
            , bgroup "selected pages index" (indexBenchmark <$> fixtures)
            , bgroup "selected pages raw" (rawBenchmark <$> fixtures)
            , bgroup "selected pages buffers" (bufferBenchmark <$> fixtures)
            ]
          else []
  defaultMain ([bgroup "selected pages" (benchmark <$> fixtures)] ++ planned)
  where
    benchmark fixture = bench (fixtureName fixture) (nf (selected (fixturePredicate fixture)) (fixtureBytes fixture))
    bufferBenchmark fixture = bench (fixtureName fixture) (nf (selectedBuffers fixture) (fixtureBytes fixture))
    plannedBenchmark fixture =
      bench
        (fixtureName fixture)
        (nf (selectedPlanned (fixturePredicate fixture) (fixtureField fixture)) (fixtureFile fixture))
    indexBenchmark fixture =
      bench
        (fixtureName fixture)
        (nf (selectedIndexRows (fixturePredicate fixture)) (fixtureFile fixture))
    rawBenchmark fixture =
      bench (fixtureName fixture) (nf (selectedRaw fixture) (fixtureBytes fixture))