Initial Implementation of the Parquet Writer

dataframe-parquet/benchmark/Writer10GB.hs

@@ -0,0 +1,49 @@
+module Main (main) where
+
+import Control.DeepSeq (NFData (rnf))
+import Criterion.Main (bench, defaultMain, envWithCleanup, whnfIO)
+import DataFrame.IO.Parquet.Writer (writeParquet)
+import DataFrame.Internal.DataFrame (DataFrame, forceDataFrame)
+import DataFrame10GB (stressDataFrame)
+import System.Directory (removeDirectoryRecursive)
+import System.FilePath ((</>))
+import System.IO.Temp (createTempDirectory, getCanonicalTemporaryDirectory)
+
+data BenchmarkEnvironment = BenchmarkEnvironment
+    { benchmarkDataFrame :: DataFrame
+    , benchmarkDirectory :: FilePath
+    , benchmarkOutput :: FilePath
+    }
+
+instance NFData BenchmarkEnvironment where
+    rnf environment =
+        forceDataFrame (benchmarkDataFrame environment) `seq`
+            rnf (benchmarkDirectory environment) `seq`
+                rnf (benchmarkOutput environment)
+
+prepareEnvironment :: IO BenchmarkEnvironment
+prepareEnvironment = do
+    temporary <- getCanonicalTemporaryDirectory
+    directory <- createTempDirectory temporary "dataframe-parquet-writer-10gb"
+    pure
+        BenchmarkEnvironment
+            { benchmarkDataFrame = stressDataFrame
+            , benchmarkDirectory = directory
+            , benchmarkOutput = directory </> "benchmark.parquet"
+            }
+
+cleanupEnvironment :: BenchmarkEnvironment -> IO ()
+cleanupEnvironment = removeDirectoryRecursive . benchmarkDirectory
+
+main :: IO ()
+main =
+    defaultMain
+        [ envWithCleanup prepareEnvironment cleanupEnvironment $ \environment ->
+            -- Memory usage for this benchmark will be north of 20 GB.
+            bench "write 10 GiB dataframe" $
+                whnfIO
+                    ( writeParquet
+                        (benchmarkOutput environment)
+                        (benchmarkDataFrame environment)
+                    )
+        ]

dataframe-parquet/dataframe-parquet.cabal

@@ -19,6 +19,7 @@ maintainer:         mschavinda@gmail.com
 copyright:          (c) 2024-2026 Michael Chavinda
 category:           Data
 tested-with:        GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2
+extra-source-files: tests/data/*.parquet
 
 common warnings
     ghc-options:
@@ -28,6 +29,11 @@ common warnings
         -Wunused-local-binds
         -Wunused-packages
 
+flag stress-tests
+    description:        Build and run the opt-in 10 GiB Parquet roundtrip stress test.
+    default:            False
+    manual:             True
+
 library
     import:             warnings
     ghc-options:        -O2
@@ -44,6 +50,11 @@ library
                         DataFrame.IO.Parquet.Thrift
                         DataFrame.IO.Parquet.Time
                         DataFrame.IO.Parquet.Utils
+                        DataFrame.IO.Parquet.Writer
+                        DataFrame.IO.Parquet.Writer.DefLevels
+                        DataFrame.IO.Parquet.Writer.Encoder
+                        DataFrame.IO.Parquet.Writer.Metadata
+                        DataFrame.IO.Parquet.Writer.Options
                         DataFrame.IO.Utils.RandomAccess
                         DataFrame.Typed.IO.Parquet
     build-depends:      base >= 4 && < 5,
@@ -52,6 +63,7 @@ library
                         dataframe-core >= 2.4.1 && < 2.5,
                         dataframe-operations >= 2.4 && < 2.5,
                         dataframe-parsing >= 2.2 && < 2.3,
+                        primitive >= 0.7 && < 0.11,
                         directory >= 1.3.0.0 && < 2,
                         filepath >= 1.4 && < 2,
                         Glob >= 0.10 && < 1,
@@ -64,3 +76,53 @@ library
                         zstd >= 0.1.2.0 && < 0.3
     hs-source-dirs:     src
     default-language:   Haskell2010
+
+
+test-suite dataframe-parquet-tests
+    import:             warnings
+    type:               exitcode-stdio-1.0
+    main-is:            Main.hs
+    hs-source-dirs:     tests
+    build-depends:      base >= 4 && < 5,
+                        bytestring >= 0.11 && < 0.14,
+                        dataframe-parquet,
+                        filepath >= 1.4 && < 2,
+                        temporary >= 1.3 && < 1.5,
+                        HUnit >= 1.6 && < 1.8
+    default-language:   Haskell2010
+
+executable dataframe-parquet-10gb-stress
+    import:             warnings
+    main-is:            StressMain.hs
+    other-modules:      DataFrame10GB
+    hs-source-dirs:     stress
+    build-depends:      base >= 4 && < 5,
+                        dataframe-core >= 2.4 && < 2.5,
+                        dataframe-parquet,
+                        filepath >= 1.4 && < 2,
+                        temporary >= 1.3 && < 1.5,
+                        text >= 2.1 && < 3,
+                        time >= 1.12 && < 2,
+                        vector >= 0.13 && < 0.15
+    default-language:   Haskell2010
+    -- ghc-options:        -O2 -threaded -rtsopts -with-rtsopts=-N
+
+benchmark dataframe-parquet-writer-10gb
+    import:             warnings
+    type:               exitcode-stdio-1.0
+    main-is:            Writer10GB.hs
+    other-modules:      DataFrame10GB
+    hs-source-dirs:     benchmark, stress
+    build-depends:      base >= 4 && < 5,
+                        criterion >= 1 && < 2,
+                        deepseq >= 1.4 && < 2,
+                        dataframe-core >= 2.4 && < 2.5,
+                        dataframe-parquet,
+                        directory >= 1.3 && < 2,
+                        filepath >= 1.4 && < 2,
+                        temporary >= 1.3 && < 1.5,
+                        text >= 2.1 && < 3,
+                        time >= 1.12 && < 2,
+                        vector >= 0.13 && < 0.15
+    default-language:   Haskell2010
+    ghc-options:        -O2 -threaded -rtsopts -with-rtsopts=-N