{-# LANGUAGE DataKinds #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeApplications #-}
-- | Tests for the writer-buffer logic in "DataFrame.IO.Utils.RandomAccess".
module Main where
import Control.Exception (SomeException, catch, evaluate)
import qualified Data.ByteString as BS
import Data.List (sortOn)
import qualified System.Exit as Exit
import System.FilePath ((</>))
import System.IO.Temp (withSystemTempDirectory)
import Test.HUnit
import Control.Monad (unless)
import Control.Monad.ST (stToIO)
import Data.Int (Int32, Int64)
import Data.Maybe (fromJust)
import qualified Data.Text as T
import DataFrame.IO.Parquet (
ParquetReadOptions (..),
defaultParquetReadOptions,
readParquet,
readParquetFiles,
readParquetWithOpts,
_readParquetWithOpts,
)
import DataFrame.IO.Parquet.Writer (
ParquetWriteOptions (..),
defaultParquetWriteOptions,
writeParquet,
writeParquetWithOptions,
)
import DataFrame.IO.Utils.RandomAccess
import DataFrame.Internal.Column (columnTypeString, fromList)
import DataFrame.Internal.DataFrame (
DataFrame,
columnNames,
fromColumns,
getColumn,
)
import DataFrame.Operations.Subset (range, select)
import qualified DataFrame.Typed.IO.Parquet as Typed
import qualified DataFrame.Typed.Operations as DT
import System.Directory (listDirectory)
directWrites :: Test
directWrites = TestCase $ do
buffer <- stToIO (mallocBuffer 1)
stToIO (writeWord8 buffer 0xaa)
stToIO (writeWord32LE buffer 0x78563412)
stToIO (writeWord64LE buffer 0x0807060504030201)
stToIO (writeFloatLE buffer 1)
stToIO (writeDoubleLE buffer 1)
writeByteString buffer (BS.pack [0xfe, 0xff])
residency <- stToIO (bufferResidency buffer)
bytes <- bufferToByteString buffer
assertEqual "direct write residency" 27 residency
assertEqual
"direct write bytes"
( BS.pack
[ 0xaa
, 0x12
, 0x34
, 0x56
, 0x78
, 0x01
, 0x02
, 0x03
, 0x04
, 0x05
, 0x06
, 0x07
, 0x08
, 0x00
, 0x00
, 0x80
, 0x3f
, 0x00
, 0x00
, 0x00
, 0x00
, 0x00
, 0x00
, 0xf0
, 0x3f
, 0xfe
, 0xff
]
)
bytes
directBufferFlush :: Test
directBufferFlush = TestCase $ do
source <- stToIO (mallocBuffer 0)
destination <- stToIO (mallocBuffer 0)
writeByteString destination (BS.pack [1, 2])
writeByteString source (BS.pack [3, 4, 5])
stToIO (flushBufferToBuffer source destination)
sourceResidency <- stToIO (bufferResidency source)
destinationBytes <- bufferToByteString destination
assertEqual "source cleared" 0 sourceResidency
assertEqual "destination appended" (BS.pack [1, 2, 3, 4, 5]) destinationBytes
stToIO (flushBufferToBuffer destination destination)
selfFlushedBytes <- bufferToByteString destination
assertEqual "self flush is a no-op" destinationBytes selfFlushedBytes
stToIO (resetPosition destination)
destinationResidency <- stToIO (bufferResidency destination)
assertEqual "reset position" 0 destinationResidency
directFileFlush :: Test
directFileFlush = TestCase $
withSystemTempDirectory "dfpq-buffer" $ \dir -> do
let outPath = dir </> "out.bin"
payload = BS.pack (take 300000 (cycle [0 .. 255]))
buffer <- stToIO (mallocBuffer 1)
writeByteString buffer payload
withWritableBinaryFile outPath $ \output ->
flushBufferToFile output buffer
residency <- stToIO (bufferResidency buffer)
contents <- BS.readFile outPath
assertEqual "source cleared after file flush" 0 residency
assertEqual "large payload round-trips" payload contents
writerRoundTrip :: String -> FilePath -> Test
writerRoundTrip label path = TestCase $
withSystemTempDirectory "dfpq-writer" $ \dir -> do
df <- readParquet path
let out = dir </> "out.parquet"
writeParquet out df
df' <- readParquet out
assertEqual label df df'
writerRoundTripTiny :: String -> FilePath -> Test
writerRoundTripTiny label path = TestCase $
withSystemTempDirectory "dfpq-writer" $ \dir -> do
df <- readParquet path
let out = dir </> "out.parquet"
writeParquetWithOptions tinyWriteOpts out df
df' <- readParquet out
assertEqual label df df'
-- Row groups of 512 rows give the parallel reader several tasks per column.
parallelMatchesSequential :: String -> FilePath -> Test
parallelMatchesSequential label path = TestCase $
withSystemTempDirectory "dfpq-parallel" $ \dir -> do
df <- readParquet path
let out = dir </> "out.parquet"
writeParquetWithOptions tinyWriteOpts out df
sequential <-
readParquetWithOpts defaultParquetReadOptions{parallel = False} out
parallelRead <-
readParquetWithOpts defaultParquetReadOptions{parallel = True} out
assertEqual (label ++ ": sequential") df sequential
assertEqual (label ++ ": parallel") df parallelRead
writerRoundTripLargeText :: Test
writerRoundTripLargeText = TestCase $
withSystemTempDirectory "dfpq-writer" $ \dir -> do
let payload = T.replicate 4096 "abcdefgh"
df = fromColumns [("text", fromList [payload, "short"])]
firstOut = dir </> "large-text-1.parquet"
secondOut = dir </> "large-text-2.parquet"
writeParquetWithOptions tinyWriteOpts firstOut df
firstRoundTrip <- readParquet firstOut
writeParquetWithOptions tinyWriteOpts secondOut firstRoundTrip
secondRoundTrip <- readParquet secondOut
assertEqual "large text first round-trip" df firstRoundTrip
assertEqual "large text second round-trip" df secondRoundTrip
-- | Sharded writes split the frame and reading them back matches.
writerRoundTripSharded :: String -> FilePath -> Int -> Int -> Test
writerRoundTripSharded label path rowsPerFile expectedShards = TestCase $
withSystemTempDirectory "dfpq-writer" $ \dir -> do
df <- readParquet path
let pattern_ = dir </> "shards" </> "part-*.parquet"
writeParquetWithOptions
defaultParquetWriteOptions{maxRowsPerFile = Just rowsPerFile}
pattern_
df
shards <- listDirectory (dir </> "shards")
assertEqual (label <> ": shard count") expectedShards (Prelude.length shards)
assertEqual
(label <> ": shard names")
["part-" <> pad i <> ".parquet" | i <- [0 .. expectedShards - 1]]
(sortOn id shards)
df' <- readParquetFiles pattern_
assertEqual (label <> ": shards round-trip") df df'
where
pad i = let s = show i in replicate (5 - Prelude.length s) '0' <> s
-- | A path without a @*@ placeholder is rejected when sharding is requested.
shardedWriteRequiresPattern :: Test
shardedWriteRequiresPattern = TestCase $
withSystemTempDirectory "dfpq-writer" $ \dir -> do
df <- readParquet "tests/data/mtcars.parquet"
threw <-
( False
<$ writeParquetWithOptions
defaultParquetWriteOptions{maxRowsPerFile = Just 4}
(dir </> "out.parquet")
df
)
`catch` (\e -> True <$ evaluate (Prelude.length (show (e :: SomeException))))
unless threw (assertFailure "expected an error for a path without '*'")
-- | Int, Int64 and Integer columns round-trip with their Haskell type.
writerRoundTripNativeIntTypes :: Test
writerRoundTripNativeIntTypes = TestCase $
withSystemTempDirectory "dfpq-writer" $ \dir -> do
let df =
fromColumns
[ ("int", fromList [1 :: Int, 2, 3])
, ("int64", fromList [1 :: Int64, 2, 3])
, ("int32", fromList [1 :: Int32, 2, 3])
, ("integer", fromList [1 :: Integer, 2, 3])
, ("nullableInt", fromList [Just (1 :: Int), Nothing, Just 3])
, ("nullableInt64", fromList [Just (1 :: Int64), Nothing, Just 3])
]
out = dir </> "int-types.parquet"
writeParquet out df
df' <- readParquet out
assertEqual
"native int types: column types"
(columnTypes df)
(columnTypes df')
assertEqual "native int types: frame" df df'
columnTypes :: DataFrame -> [(String, String)]
columnTypes df =
[ (T.unpack name, columnTypeString (fromJust (getColumn name df)))
| name <- columnNames df
]
tinyWriteOpts :: ParquetWriteOptions
tinyWriteOpts =
defaultParquetWriteOptions
{ pageSize = 64
, rowGroupSize = 512
, batchRows = 4
, subBatchRows = 3
}
dottedColumnNames :: Test
dottedColumnNames = TestCase $
withSystemTempDirectory "dfpq-dotted" $ \dir -> do
let df =
fromColumns
[ ("sepal.length", fromList [5.1, 4.9, 4.7 :: Double])
, ("length", fromList [1, 2, 3 :: Int])
, ("variety", fromList ["a", "b", "c" :: T.Text])
]
out = dir </> "dotted.parquet"
writeParquet out df
whole <- readParquet out
assertEqual "names survive the round trip" (columnNames df) (columnNames whole)
projected <-
readParquetWithOpts
defaultParquetReadOptions{selectedColumns = Just ["sepal.length", "variety"]}
out
assertEqual
"projection keeps the dotted column"
["sepal.length", "variety"]
(columnNames projected)
assertEqual
"projected rows match"
(select ["sepal.length", "variety"] df)
projected
leaf <-
readParquetWithOpts
defaultParquetReadOptions{selectedColumns = Just ["length"]}
out
assertEqual "a full name wins over a leaf alias" ["length"] (columnNames leaf)
typed <-
Typed.readParquet @'[ '("sepal.length", Double), '("variety", T.Text)] out
assertEqual "the typed reader projects a dotted schema" 3 (DT.nRows typed)
tests :: Test
tests =
TestList
[ TestLabel "dotted column names project by full name" dottedColumnNames
, TestLabel "direct buffer writes" directWrites
, TestLabel "direct buffer-to-buffer flush" directBufferFlush
, TestLabel "direct buffer-to-file flush" directFileFlush
, TestLabel
"writer roundtrip: alltypes_plain"
(writerRoundTrip "alltypes_plain" "tests/data/alltypes_plain.parquet")
, TestLabel
"writer roundtrip: alltypes_plain.snappy"
( writerRoundTrip
"alltypes_plain.snappy"
"tests/data/alltypes_plain.snappy.parquet"
)
, TestLabel
"writer roundtrip: alltypes_dictionary"
(writerRoundTrip "alltypes_dictionary" "tests/data/alltypes_dictionary.parquet")
, TestLabel
"writer roundtrip: alltypes_tiny_pages"
(writerRoundTrip "alltypes_tiny_pages" "tests/data/alltypes_tiny_pages.parquet")
, TestLabel
"parallel read matches sequential: transactions"
(parallelMatchesSequential "transactions" "tests/data/transactions.parquet")
, TestLabel
"parallel read matches sequential: mtcars"
(parallelMatchesSequential "mtcars" "tests/data/mtcars.parquet")
, TestLabel
"writer roundtrip: transactions"
(writerRoundTrip "transactions" "tests/data/transactions.parquet")
, TestLabel
"writer roundtrip: mtcars"
(writerRoundTrip "mtcars" "tests/data/mtcars.parquet")
, TestLabel
"writer roundtrip: int32_decimal"
(writerRoundTrip "int32_decimal" "tests/data/int32_decimal.parquet")
, TestLabel
"writer roundtrip: int64_decimal"
(writerRoundTrip "int64_decimal" "tests/data/int64_decimal.parquet")
, TestLabel
"writer roundtrip: sharded mtcars"
(writerRoundTripSharded "sharded mtcars" "tests/data/mtcars.parquet" 10 4)
, TestLabel
"writer roundtrip: sharded exact multiple"
(writerRoundTripSharded "sharded exact" "tests/data/mtcars.parquet" 32 1)
, TestLabel
"sharded write requires a '*' pattern"
shardedWriteRequiresPattern
, TestLabel
"writer roundtrip: Int/Integer keep their Haskell type"
writerRoundTripNativeIntTypes
, TestLabel
"writer roundtrip: alltypes_plain multi-page"
( writerRoundTripTiny
"alltypes_plain multi-page"
"tests/data/alltypes_plain.parquet"
)
, TestLabel "writer roundtrip: large text" writerRoundTripLargeText
, TestLabel "parallel read: many row groups" parallelReadRowGroups
, TestLabel "parallel read: row range skips row groups" parallelReadRowRange
, TestLabel "parallel read: widened integers" parallelReadWidenIntegers
]
-- | A frame big enough for 'tinyWriteOpts' to split into many row groups.
manyRowGroupsFrame :: DataFrame
manyRowGroupsFrame =
fromColumns
[ ("n", fromList [i * 7 - 3000 | i <- [0 .. 2999 :: Int]])
, ("d", fromList [fromIntegral i / 8 :: Double | i <- [0 .. 2999 :: Int]])
,
( "t"
, fromList
[T.pack (show i) <> T.replicate (i `mod` 5) "é" | i <- [0 .. 2999 :: Int]]
)
]
withManyRowGroups :: (FilePath -> IO ()) -> IO ()
withManyRowGroups body = withSystemTempDirectory "dfpq-rg" $ \dir -> do
let out = dir </> "rg.parquet"
writeParquetWithOptions tinyWriteOpts out manyRowGroupsFrame
body out
-- | The parallel reader matches the sequential one and the source frame.
parallelReadRowGroups :: Test
parallelReadRowGroups = TestCase $ withManyRowGroups $ \out -> do
parallel <- readParquet out
sequential <- _readParquetWithOpts (Just True) defaultParquetReadOptions out
assertEqual "parallel vs source" manyRowGroupsFrame parallel
assertEqual "parallel vs sequential" sequential parallel
projected <-
readParquetWithOpts
defaultParquetReadOptions{selectedColumns = Just ["t", "n"]}
out
assertEqual "projection" (select ["t", "n"] manyRowGroupsFrame) projected
-- | A row range only decodes the row groups it overlaps, and slices exactly.
parallelReadRowRange :: Test
parallelReadRowRange = TestCase $ withManyRowGroups $ \out -> do
let ranges = [(0, 5), (7, 1300), (1299, 1301), (2990, 3000), (2995, 4000), (100, 100)]
mapM_
( \r -> do
got <- readParquetWithOpts defaultParquetReadOptions{rowRange = Just r} out
assertEqual ("range " ++ show r) (range r manyRowGroupsFrame) got
)
ranges
-- | 'widenIntegers' turns Int32 and Int64 columns into Int on both read paths.
parallelReadWidenIntegers :: Test
parallelReadWidenIntegers = TestCase $ withSystemTempDirectory "dfpq-widen" $ \dir -> do
let n = 3000 :: Int
df =
fromColumns
[ ("i32", fromList [fromIntegral (i - 1500) :: Int32 | i <- [0 .. n - 1]])
, ("i64", fromList [fromIntegral i * 4000000000 :: Int64 | i <- [0 .. n - 1]])
]
expected =
fromColumns
[ ("i32", fromList [i - 1500 | i <- [0 .. n - 1]])
, ("i64", fromList [i * 4000000000 | i <- [0 .. n - 1]])
]
out = dir </> "ints.parquet"
opts = defaultParquetReadOptions{widenIntegers = True}
writeParquetWithOptions tinyWriteOpts out df
parallel <- readParquetWithOpts opts out
sequential <- _readParquetWithOpts (Just True) opts out
assertEqual "parallel" expected parallel
assertEqual "sequential" expected sequential
main :: IO ()
main = do
result <- runTestTT tests
if failures result > 0 || errors result > 0
then Exit.exitFailure
else Exit.exitSuccess