diff --git a/dataframe-learn/src-internal/DataFrame/DecisionTree/Cart.hs b/dataframe-learn/src-internal/DataFrame/DecisionTree/Cart.hs index e3bbc0c4..30aed2dd 100644 --- a/dataframe-learn/src-internal/DataFrame/DecisionTree/Cart.hs +++ b/dataframe-learn/src-internal/DataFrame/DecisionTree/Cart.hs @@ -344,7 +344,8 @@ nullableLeq c | otherwise = Nothing oneHotFeatures :: - forall b. (Columnable b) => T.Text -> Maybe Bitmap -> V.Vector b -> [CartFeature] + forall b. + (Columnable b) => T.Text -> Maybe Bitmap -> V.Vector b -> [CartFeature] oneHotFeatures c bm v = case testEquality (typeRep @b) (typeRep @T.Text) of Just Refl -> [oneHot c nulls v cat | cat <- Set.toList (Set.fromList present)] Nothing -> [] @@ -352,7 +353,8 @@ oneHotFeatures c bm v = case testEquality (typeRep @b) (typeRep @T.Text) of nulls = fmap (nullFlags (V.length v)) bm present = [x | (i, x) <- zip [0 ..] (V.toList v), maybe True (not . (VU.! i)) nulls] -oneHot :: T.Text -> Maybe (VU.Vector Bool) -> V.Vector T.Text -> T.Text -> CartFeature +oneHot :: + T.Text -> Maybe (VU.Vector Bool) -> V.Vector T.Text -> T.Text -> CartFeature oneHot c Nothing v cat = CartFeature (VU.generate (V.length v) (\i -> if v V.! i == cat then 1 else 0)) diff --git a/dataframe-learn/src-internal/DataFrame/DecisionTree/Histogram.hs b/dataframe-learn/src-internal/DataFrame/DecisionTree/Histogram.hs index d4c78803..f1c15a87 100644 --- a/dataframe-learn/src-internal/DataFrame/DecisionTree/Histogram.hs +++ b/dataframe-learn/src-internal/DataFrame/DecisionTree/Histogram.hs @@ -54,7 +54,8 @@ binFeature maxBins f = Binned f bins nb upper runs = runLengths sorted distinct = map fst runs edges - | length runs <= maxBins = VU.fromList (zipWith splitMidpoint distinct (drop 1 distinct)) + | length runs <= maxBins = + VU.fromList (zipWith splitMidpoint distinct (drop 1 distinct)) | otherwise = VU.fromList (cutEdges maxBins (VU.length sorted) runs) nb = VU.length edges + 1 -- The last bin's threshold is the largest value, so "every non-null row @@ -112,7 +113,8 @@ Built strictly: left lazy, the histograms of a deep tree pile up as thunks. -} type Hist = VU.Vector Double -buildHist :: VU.Vector Double -> VU.Vector Double -> VU.Vector Int -> Binned -> Hist +buildHist :: + VU.Vector Double -> VU.Vector Double -> VU.Vector Int -> Binned -> Hist buildHist w y idxs bn = runST $ do m <- VUM.replicate (4 * (bnNBins bn + 1)) 0 VU.forM_ idxs $ \i -> do @@ -141,7 +143,7 @@ fitBinnedTree :: fitBinnedTree lim binned y mw = (toTree root, inSample) where n = VU.length y - w = maybe (VU.replicate n 1) id mw + w = Data.Maybe.fromMaybe (VU.replicate n 1) mw allIdx = VU.enumFromN 0 n root = node 0 allIdx (hists allIdx) inSample = VU.update (VU.replicate n 0) (VU.concat (leafRows root)) @@ -155,7 +157,8 @@ fitBinnedTree lim binned y mw = (toTree root, inSample) strategy = if VU.length idxs >= 20000 then parList rseq else evalList rseq node depth idxs hs - | depth >= tlMaxDepth lim || VU.length idxs < tlMinSamplesSplit lim || V.null hs = leaf + | depth >= tlMaxDepth lim || VU.length idxs < tlMinSamplesSplit lim || V.null hs = + leaf | otherwise = maybe leaf split (bestSplit lim binned hs) where Totals tw tsy _ _ = histTotals (V.head hs) @@ -199,7 +202,8 @@ histTotals h = go 0 0 0 0 0 scored with the node's null rows on the right and then on the left; ties keep the earliest candidate, as in the exact sweep. -} -bestSplit :: TreeLimits -> V.Vector Binned -> V.Vector Hist -> Maybe (Int, Int, Bool) +bestSplit :: + TreeLimits -> V.Vector Binned -> V.Vector Hist -> Maybe (Int, Int, Bool) bestSplit lim binned hs | null candidates = Nothing | red > 0 && red >= tlMinImpurityDecrease lim = Just sp @@ -225,7 +229,14 @@ bestSplit lim binned hs | otherwise = Nothing where wr = totW - wl - scan :: Int -> Double -> Double -> Double -> Int -> Maybe (Int, Bool, Double) -> Maybe (Int, Bool, Double) + scan :: + Int -> + Double -> + Double -> + Double -> + Int -> + Maybe (Int, Bool, Double) -> + Maybe (Int, Bool, Double) scan !b !wl !syl !syl2 !cl best | b > lastB = best | otherwise = scan (b + 1) wl' syl' syl2' cl' best' @@ -236,7 +247,8 @@ bestSplit lim binned hs cl' = cl + round (h VU.! (4 * b + 3)) right = (,) False <$> score cl' wl' syl' syl2' left - | hasNulls = (,) True <$> score (cl' + nC) (wl' + nW) (syl' + nSY) (syl2' + nSY2) + | hasNulls = + (,) True <$> score (cl' + nC) (wl' + nW) (syl' + nSY) (syl2' + nSY2) | otherwise = Nothing consider bst (dir, r) | maybe True (\(_, _, rb) -> r > rb) bst = Just (b, dir, r) diff --git a/dataframe-learn/src/DataFrame/Boosting/AdaBoost.hs b/dataframe-learn/src/DataFrame/Boosting/AdaBoost.hs index 130b774b..2760f16c 100644 --- a/dataframe-learn/src/DataFrame/Boosting/AdaBoost.hs +++ b/dataframe-learn/src/DataFrame/Boosting/AdaBoost.hs @@ -30,10 +30,10 @@ import DataFrame.Errors (DataFrameException (..)) import DataFrame.DecisionTree.Cart ( CartFeature (..), - cfPred, - splitMidpoint, cartFeatures, + cfPred, sortIndicesByValue, + splitMidpoint, ) import DataFrame.DecisionTree.Fit (treeToExpr) import DataFrame.DecisionTree.Types (Tree (..)) diff --git a/dataframe-learn/src/DataFrame/Boosting/GBM.hs b/dataframe-learn/src/DataFrame/Boosting/GBM.hs index bc5c5900..e59bab99 100644 --- a/dataframe-learn/src/DataFrame/Boosting/GBM.hs +++ b/dataframe-learn/src/DataFrame/Boosting/GBM.hs @@ -3,7 +3,7 @@ {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} + {-# LANGUAGE TypeFamilies #-} {- | Gradient boosting of regression trees (Friedman). Trees are fitted to the @@ -34,7 +34,11 @@ import DataFrame.Errors (DataFrameException (..)) import DataFrame.DecisionTree.Cart (cartFeatures) import DataFrame.DecisionTree.Fit (treeToExpr) -import DataFrame.DecisionTree.Histogram (TreeLimits (..), binFeatures, fitBinnedTree) +import DataFrame.DecisionTree.Histogram ( + TreeLimits (..), + binFeatures, + fitBinnedTree, + ) import DataFrame.DecisionTree.Types (Tree) import DataFrame.Expression.Operators ((.*.), (.+.), (.>.)) import DataFrame.Featurize.Internal (targetDoubles) @@ -53,8 +57,9 @@ data GBConfig = GBConfig , gbLearningRate :: !Double , gbMaxDepth :: !Int , gbMaxBins :: !Int - -- ^ Bins per feature for split finding. Features with at most this many - -- distinct values split exactly as on the raw values. + {- ^ Bins per feature for split finding. Features with at most this many + distinct values split exactly as on the raw values. + -} , gbSeed :: !Int } deriving (Eq, Show) diff --git a/dataframe-learn/src/DataFrame/DecisionTree/Model.hs b/dataframe-learn/src/DataFrame/DecisionTree/Model.hs index 79b2240a..fba2a577 100644 --- a/dataframe-learn/src/DataFrame/DecisionTree/Model.hs +++ b/dataframe-learn/src/DataFrame/DecisionTree/Model.hs @@ -27,7 +27,11 @@ import qualified Data.Vector as V import DataFrame.DecisionTree.Cart (cartFeatures) import DataFrame.DecisionTree.Fit (fitDecisionTree, treeToExpr) -import DataFrame.DecisionTree.Regression (RegFit (..), RegTreeConfig, fitRegTree) +import DataFrame.DecisionTree.Regression ( + RegFit (..), + RegTreeConfig, + fitRegTree, + ) import DataFrame.DecisionTree.Types (Tree (..), TreeConfig) import DataFrame.Featurize.Internal (targetDoubles) import DataFrame.Internal.Column (Columnable) diff --git a/dataframe-learn/src/DataFrame/DecisionTree/Regression.hs b/dataframe-learn/src/DataFrame/DecisionTree/Regression.hs index 6995378e..b4f2b34b 100644 --- a/dataframe-learn/src/DataFrame/DecisionTree/Regression.hs +++ b/dataframe-learn/src/DataFrame/DecisionTree/Regression.hs @@ -83,8 +83,8 @@ fitRegTree cfg feats y mw = RegFit (toTree root) inSample splitNode depth idxs sortedByFeat (fj, thr, side) | VU.null lefts || VU.null rights = FLeaf (weightedMean idxs) idxs | otherwise = - forceTree l - `par` (forceTree r `pseq` FBranch (cfSplit feat thr side) l r) + forceTree l `par` + (forceTree r `pseq` FBranch (cfSplit feat thr side) l r) where feat = feats V.! fj vals = cfValues feat @@ -134,7 +134,8 @@ fitRegTree cfg feats y mw = RegFit (toTree root) inSample lastK = if hasNulls then m - 1 else m - 2 score nl wl syl syl2 = let wr = totW - wl - ok = nl >= rtMinLeafSize cfg && nNode - nl >= rtMinLeafSize cfg && wl > 0 && wr > 0 + ok = + nl >= rtMinLeafSize cfg && nNode - nl >= rtMinLeafSize cfg && wl > 0 && wr > 0 in if ok then Just (nodeSSE - (sse syl syl2 wl + sse (totSY - syl) (totSY2 - syl2) wr)) else Nothing diff --git a/dataframe-learn/tests-internal/HistogramTrees.hs b/dataframe-learn/tests-internal/HistogramTrees.hs index 5422567e..ae315265 100644 --- a/dataframe-learn/tests-internal/HistogramTrees.hs +++ b/dataframe-learn/tests-internal/HistogramTrees.hs @@ -16,7 +16,11 @@ import Test.HUnit import DataFrame.DecisionTree.Cart (cartFeatures) import DataFrame.DecisionTree.Fit (treeToExpr) -import DataFrame.DecisionTree.Histogram (TreeLimits (..), binFeatures, fitBinnedTree) +import DataFrame.DecisionTree.Histogram ( + TreeLimits (..), + binFeatures, + fitBinnedTree, + ) import DataFrame.DecisionTree.Regression ( RegFit (..), RegTreeConfig (..), @@ -30,17 +34,26 @@ import qualified DataFrameApi as D tests :: [Test] tests = - [ TestLabel "histogram tree: in-sample == interpreted (mixed nullable frame)" inSampleMatchesInterpret - , TestLabel "histogram tree: in-sample == interpreted (more values than bins)" coarseBinsMatchInterpret + [ TestLabel + "histogram tree: in-sample == interpreted (mixed nullable frame)" + inSampleMatchesInterpret + , TestLabel + "histogram tree: in-sample == interpreted (more values than bins)" + coarseBinsMatchInterpret , TestLabel "histogram tree: same partition as the exact tree" matchesExactTree ] limits :: Int -> TreeLimits limits depth = TreeLimits depth 2 1 0 -fitBinned :: Int -> Int -> D.DataFrame -> VU.Vector Double -> (Tree Double, VU.Vector Double) +fitBinned :: + Int -> Int -> D.DataFrame -> VU.Vector Double -> (Tree Double, VU.Vector Double) fitBinned bins depth df y = - fitBinnedTree (limits depth) (binFeatures bins (V.fromList (cartFeatures "y" df))) y Nothing + fitBinnedTree + (limits depth) + (binFeatures bins (V.fromList (cartFeatures "y" df))) + y + Nothing interpreted :: D.DataFrame -> Tree Double -> VU.Vector Double interpreted df t = case interpret @Double df (treeToExpr t) of @@ -55,18 +68,37 @@ mixedFrame :: Int -> ([Double], D.DataFrame) mixedFrame n = (y, D.fromColumns (("y", DI.fromList y) : cols)) where rows = [0 .. n - 1] - md = [if i `mod` 7 == 0 then Nothing else Just (fromIntegral ((i * 37) `mod` 23) :: Double) | i <- rows] - mi = [if i `mod` 5 == 2 then Nothing else Just ((i * 11) `mod` 9 :: Int) | i <- rows] - mt = [if i `mod` 6 == 4 then Nothing else Just (["a", "b", "c"] !! (i `mod` 3) :: T.Text) | i <- rows] + md = + [ if i `mod` 7 == 0 + then Nothing + else Just (fromIntegral ((i * 37) `mod` 23) :: Double) + | i <- rows + ] + mi = + [if i `mod` 5 == 2 then Nothing else Just ((i * 11) `mod` 9 :: Int) | i <- rows] + mt = + [ if i `mod` 6 == 4 + then Nothing + else Just (["a", "b", "c"] !! (i `mod` 3) :: T.Text) + | i <- rows + ] z = [fromIntegral ((i * 13) `mod` 17) :: Double | i <- rows] y = [fromIntegral ((i * 29) `mod` 31) / 7 :: Double | i <- rows] - cols = [("md", maybeCol md), ("mi", maybeCol mi), ("mt", maybeCol mt), ("z", DI.fromList z)] + cols = + [ ("md", maybeCol md) + , ("mi", maybeCol mi) + , ("mt", maybeCol mt) + , ("z", DI.fromList z) + ] inSampleMatchesInterpret :: Test inSampleMatchesInterpret = TestCase $ do let (y, df) = mixedFrame 60 (t, inSample) = fitBinned 1024 6 df (VU.fromList y) - assertEqual "in-sample predictions equal the interpreted tree" (interpreted df t) inSample + assertEqual + "in-sample predictions equal the interpreted tree" + (interpreted df t) + inSample -- 23 distinct values of md squeezed into 4 bins: thresholds come from the -- equal-count cuts, and routing must still agree. @@ -74,7 +106,10 @@ coarseBinsMatchInterpret :: Test coarseBinsMatchInterpret = TestCase $ do let (y, df) = mixedFrame 200 (t, inSample) = fitBinned 4 5 df (VU.fromList y) - assertEqual "in-sample predictions equal the interpreted tree" (interpreted df t) inSample + assertEqual + "in-sample predictions equal the interpreted tree" + (interpreted df t) + inSample matchesExactTree :: Test matchesExactTree = TestCase $ do diff --git a/dataframe-learn/tests-internal/NullSplits.hs b/dataframe-learn/tests-internal/NullSplits.hs index 93b7a35f..c46d9c35 100644 --- a/dataframe-learn/tests-internal/NullSplits.hs +++ b/dataframe-learn/tests-internal/NullSplits.hs @@ -23,15 +23,17 @@ import DataFrame.DecisionTree.Regression ( fitRegTree, ) import DataFrame.DecisionTree.Types (Tree, TreeConfig (..), defaultTreeConfig) -import DataFrame.Internal.Expression (Expr (..)) import qualified DataFrame.Internal.Column as DI +import DataFrame.Internal.Expression (Expr (..)) import DataFrame.Internal.Interpreter (interpret) import DataFrame.Model (Fit (..), Predict (..)) import qualified DataFrameApi as D tests :: [Test] tests = - [ TestLabel "null splits: in-sample == interpreted (mixed nullable frame)" inSampleMatchesInterpret + [ TestLabel + "null splits: in-sample == interpreted (mixed nullable frame)" + inSampleMatchesInterpret , TestLabel "null splits: nulls join the side they fit" learnedDirection , TestLabel "null splits: null-indicator split" nullIndicator , TestLabel "null splits: unseen nulls still route" unseenNullsRoute @@ -39,7 +41,8 @@ tests = , TestLabel "null splits: AdaBoost null indicator" adaBoostNullIndicator ] -fitAll :: Int -> D.DataFrame -> VU.Vector Double -> (Tree Double, VU.Vector Double) +fitAll :: + Int -> D.DataFrame -> VU.Vector Double -> (Tree Double, VU.Vector Double) fitAll depth df y = (rfTree fitted, rfFitted fitted) where fitted = @@ -67,17 +70,35 @@ inSampleMatchesInterpret :: Test inSampleMatchesInterpret = TestCase $ do let n = 60 :: Int rows = [0 .. n - 1] - md = [if i `mod` 7 == 0 then Nothing else Just (fromIntegral ((i * 37) `mod` 23) :: Double) | i <- rows] - mi = [if i `mod` 5 == 2 then Nothing else Just ((i * 11) `mod` 9 :: Int) | i <- rows] - mt = [if i `mod` 6 == 4 then Nothing else Just (["a", "b", "c"] !! (i `mod` 3) :: T.Text) | i <- rows] + md = + [ if i `mod` 7 == 0 + then Nothing + else Just (fromIntegral ((i * 37) `mod` 23) :: Double) + | i <- rows + ] + mi = + [if i `mod` 5 == 2 then Nothing else Just ((i * 11) `mod` 9 :: Int) | i <- rows] + mt = + [ if i `mod` 6 == 4 + then Nothing + else Just (["a", "b", "c"] !! (i `mod` 3) :: T.Text) + | i <- rows + ] z = [fromIntegral ((i * 13) `mod` 17) :: Double | i <- rows] y = [fromIntegral ((i * 29) `mod` 31) / 7 :: Double | i <- rows] df = withTarget y - [("md", maybeCol md), ("mi", maybeCol mi), ("mt", maybeCol mt), ("z", DI.fromList z)] + [ ("md", maybeCol md) + , ("mi", maybeCol mi) + , ("mt", maybeCol mt) + , ("z", DI.fromList z) + ] (t, inSample) = fitAll 6 df (VU.fromList y) - assertEqual "in-sample predictions equal the interpreted tree" (interpreted df t) inSample + assertEqual + "in-sample predictions equal the interpreted tree" + (interpreted df t) + inSample -- The null rows have the same target as the low values of x, so one split -- fits exactly only if it sends the nulls left. @@ -108,7 +129,9 @@ unseenNullsRoute = TestCase $ do scored = interpreted testDf t assertEqual "fits exactly" (VU.fromList y) inSample assertEqual "present rows unchanged" (VU.take 4 inSample) (VU.take 4 scored) - assertBool "null rows land in a leaf" (VU.all (`elem` [0, 1]) (VU.drop 4 scored)) + assertBool + "null rows land in a leaf" + (VU.all (`elem` [0, 1]) (VU.drop 4 scored)) -- The target depends only on whether x is null. If the split's threshold were -- +Infinity instead of 3, CART and AdaBoost would send the null rows left @@ -122,7 +145,11 @@ indicatorFrame = (y, withTarget y [("x", maybeCol x)]) cartNullIndicator :: Test cartNullIndicator = TestCase $ do let (y, df) = indicatorFrame - t = buildCartTree @Double defaultTreeConfig{maxTreeDepth = 1, minLeafSize = 1} "y" df + t = + buildCartTree @Double + defaultTreeConfig{maxTreeDepth = 1, minLeafSize = 1} + "y" + df assertEqual "fits exactly" (VU.fromList y) (interpreted df t) adaBoostNullIndicator :: Test diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs index 8dca6ca7..1e38029e 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer.hs @@ -1,5 +1,4 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedRecordDot #-} {-# LANGUAGE OverloadedStrings #-} @@ -13,20 +12,15 @@ module DataFrame.IO.Parquet.Writer ( nativeTypeKeyValues, ) where +import Control.Exception (evaluate) import Control.Monad (forM_, unless, when) -import Control.Monad.IO.Class (MonadIO) -import Control.Monad.Primitive (PrimBase, PrimMonad, PrimState) +import Control.Monad.ST (RealWorld, ST, stToIO) import qualified Data.ByteString as BS +import Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef) import Data.Int (Int64) import Data.Maybe (fromJust) import Data.Primitive.ByteArray (getSizeofMutableByteArray) -import Data.Primitive.MutVar ( - MutVar, - modifyMutVar', - newMutVar, - readMutVar, - writeMutVar, - ) +import Data.STRef (STRef, modifySTRef', newSTRef, readSTRef, writeSTRef) import qualified Data.Text as T import qualified Data.Vector as VB import DataFrame.IO.Parquet.Thrift hiding (schema) @@ -56,12 +50,12 @@ import DataFrame.IO.Utils.RandomAccess ( WritableBinaryHandle, atomicallyWriteFile, bufferResidency, - bufferToByteString, ensureCapacity, flushBufferToBuffer, flushBufferToFile, mallocBuffer, resetPosition, + unsafeBufferView, withWritableBinaryFile, writeByteString, writeByteStringToFile, @@ -80,29 +74,29 @@ import System.Directory (createDirectoryIfMissing) import System.FilePath (takeDirectory) import Text.Printf (printf) -data ParquetWriterState m = ParquetWriterState +data ParquetWriterState = ParquetWriterState { outputFileHandle :: !WritableBinaryHandle - , columnChunks :: !(VB.Vector (ColumnChunkState m)) - , currentFileOffsetRef :: !(MutVar (PrimState m) Int64) - , scratchBuffer :: !(MemoryBuffer (PrimState m)) - , rowGroupMetadataRef :: !(MutVar (PrimState m) [RowGroup]) - , rowNumberRef :: !(MutVar (PrimState m) Int) + , columnChunks :: !(VB.Vector (ColumnChunkState RealWorld)) + , currentFileOffsetRef :: !(IORef Int64) + , scratchBuffer :: !(MemoryBuffer RealWorld) + , rowGroupMetadataRef :: !(IORef [RowGroup]) + , rowNumberRef :: !(IORef Int) } -data ColumnChunkState m = ColumnChunkState +data ColumnChunkState s = ColumnChunkState { columnName :: !T.Text , nullable :: !Bool , schema :: !SchemaElement - , encoder :: !(Encoder m) - , buffer :: !(MemoryBuffer (PrimState m)) - , uncompressedBufferSize :: !(MutVar (PrimState m) Int64) - , pageState :: !(PageState m) + , encoder :: !(Encoder s) + , buffer :: !(MemoryBuffer s) + , uncompressedBufferSize :: !(STRef s Int64) + , pageState :: !(PageState s) } -data PageState m = PageState - { pageBuffer :: !(MemoryBuffer (PrimState m)) - , definitionLevels :: !(DefLevels (PrimState m)) - , currentRowCount :: !(MutVar (PrimState m) Int) +data PageState s = PageState + { pageBuffer :: !(MemoryBuffer s) + , definitionLevels :: !(DefLevels s) + , currentRowCount :: !(STRef s Int) } writeParquet :: FilePath -> DataFrame -> IO () @@ -163,21 +157,22 @@ writeShard options path_ df startRow endRow = do let names = columnNames df shardRows = max 0 (endRow - startRow) columnChunks_ <- - VB.fromList - <$> mapM - ( \columnName_ -> - initColumnChunkState - options - columnName_ - (fromJust (getColumn columnName_ df)) - ) - names - scratchBuffer_ <- mallocBuffer (max 1 options.pageSize) + stToIO $ + VB.fromList + <$> mapM + ( \columnName_ -> + initColumnChunkState + options + columnName_ + (fromJust (getColumn columnName_ df)) + ) + names + scratchBuffer_ <- stToIO (mallocBuffer (max 1 options.pageSize)) atomicallyWriteFile path_ $ \path -> withWritableBinaryFile path $ \output -> do writeByteStringToFile output magic - currentFileOffsetRef_ <- newMutVar 4 - rowGroupMetadataRef_ <- newMutVar [] - rowNumberRef_ <- newMutVar 0 + currentFileOffsetRef_ <- newIORef 4 + rowGroupMetadataRef_ <- newIORef [] + rowNumberRef_ <- newIORef 0 let writerState = ParquetWriterState output @@ -188,26 +183,29 @@ writeShard options path_ df startRow endRow = do rowNumberRef_ interval = max 1 options.batchRows subBatch = max 1 options.subBatchRows + writeBatch :: Int -> Int -> IO () writeBatch rowNum batchEnd | rowNum >= batchEnd = pure () | otherwise = do let count = min subBatch (batchEnd - rowNum) - VB.forM_ columnChunks_ (writeRows options scratchBuffer_ rowNum count) - modifyMutVar' rowNumberRef_ (+ count) + VB.forM_ columnChunks_ $ \ccs -> do + pageFull <- stToIO (writeRows options rowNum count ccs) + when pageFull (flushPage options scratchBuffer_ ccs) + modifyIORef' rowNumberRef_ (+ count) writeBatch (rowNum + count) batchEnd + loop :: Int -> IO () loop rowNum | rowNum >= endRow = pure () | otherwise = do let batchEnd = rowNum + min interval (endRow - rowNum) writeBatch rowNum batchEnd - size <- bufferedSize columnChunks_ + size <- stToIO (bufferedSize columnChunks_) when (size >= options.rowGroupSize) $ flushRowGroup options writerState loop batchEnd loop startRow flushRowGroup options writerState - rowGroupMetadata <- - reverse <$> readMutVar rowGroupMetadataRef_ + rowGroupMetadata <- reverse <$> readIORef rowGroupMetadataRef_ let schemaElements = rootSchemaElement (VB.length columnChunks_) : VB.toList (VB.map schema columnChunks_) @@ -229,31 +227,28 @@ nativeTypeKeyValues names df = , Just col <- [getColumn name df] ] +{- | Encode @count@ rows starting at @firstRow@ into the column's page. +Returns 'True' when the page has reached the page size and should be flushed. +-} writeRows :: - (PrimBase m, MonadIO m) => - ParquetWriteOptions -> - MemoryBuffer (PrimState m) -> - Int -> - Int -> - ColumnChunkState m -> - m () -writeRows options scratch firstRow count ccs = do + ParquetWriteOptions -> Int -> Int -> ColumnChunkState s -> ST s Bool +writeRows options firstRow count ccs = do let page = ccs.pageState buf = page.pageBuffer encode = ccs.encoder.encodeValue dl = page.definitionLevels end = firstRow + count - pos0 <- readMutVar buf.positionRef + pos0 <- readSTRef buf.positionRef let margin = options.pageSize arr0 <- ensureCapacity buf (pos0 + max margin (count * 64)) size0 <- getSizeofMutableByteArray arr0 let go !size !pos !row - | row >= end = writeMutVar buf.positionRef pos + | row >= end = writeSTRef buf.positionRef pos | pos + margin > size = do -- Rare: buffer nearly full, grow it - writeMutVar buf.positionRef pos + writeSTRef buf.positionRef pos arr' <- ensureCapacity buf @@ -269,43 +264,35 @@ writeRows options scratch firstRow count ccs = do go size0 pos0 firstRow -- Batch bookkeeping: once per sub-batch instead of per value - modifyMutVar' page.currentRowCount (+ count) + modifySTRef' page.currentRowCount (+ count) flushDef dl pageRes <- bufferResidency buf defRes <- bufferResidency dl.dlBuf - when - (pageRes + defRes >= options.pageSize) - (flushPage options scratch ccs) + pure (pageRes + defRes >= options.pageSize) flushPage :: - (PrimBase m, MonadIO m) => ParquetWriteOptions -> - MemoryBuffer (PrimState m) -> - ColumnChunkState m -> - m () + MemoryBuffer RealWorld -> + ColumnChunkState RealWorld -> + IO () flushPage options scratch columnChunkState = do let page = columnChunkState.pageState - numPageRows <- readMutVar page.currentRowCount + numPageRows <- stToIO (readSTRef page.currentRowCount) when (numPageRows > 0) $ do - pos <- readMutVar page.pageBuffer.positionRef - pos' <- columnChunkState.encoder.finishValues page.pageBuffer pos - writeMutVar page.pageBuffer.positionRef pos' - body <- assemblePageBody scratch columnChunkState - writeDataPage - options.compressionCodec - numPageRows - body - columnChunkState - resetPosition page.pageBuffer - resetPosition page.definitionLevels.dlBuf - resetPosition scratch - writeMutVar page.currentRowCount 0 + body <- stToIO $ do + pos <- readSTRef page.pageBuffer.positionRef + pos' <- columnChunkState.encoder.finishValues page.pageBuffer pos + writeSTRef page.pageBuffer.positionRef pos' + assemblePageBody scratch columnChunkState + writeDataPage options.compressionCodec numPageRows body columnChunkState + stToIO $ do + resetPosition page.pageBuffer + resetPosition page.definitionLevels.dlBuf + resetPosition scratch + writeSTRef page.currentRowCount 0 assemblePageBody :: - (PrimMonad m) => - MemoryBuffer (PrimState m) -> - ColumnChunkState m -> - m (MemoryBuffer (PrimState m)) + MemoryBuffer s -> ColumnChunkState s -> ST s (MemoryBuffer s) assemblePageBody scratch columnChunkState | not columnChunkState.nullable = pure columnChunkState.pageState.pageBuffer | otherwise = do @@ -319,18 +306,20 @@ assemblePageBody scratch columnChunkState pure scratch writeDataPage :: - (PrimBase m, MonadIO m) => CompressionCodec -> Int -> - MemoryBuffer (PrimState m) -> - ColumnChunkState m -> - m () + MemoryBuffer RealWorld -> + ColumnChunkState RealWorld -> + IO () writeDataPage codec numPageRows body columnChunkState = do - uncompressedPageSize <- bufferResidency body + uncompressedPageSize <- stToIO (bufferResidency body) compressedBody <- case codec of UNCOMPRESSED _ -> pure Nothing - SNAPPY _ -> - Just . Snappy.compress <$> bufferToByteString body + SNAPPY _ -> do + -- Compress straight from the page buffer, forcing the result now + -- so it can't be read after the buffer is reused for the next page. + page <- unsafeBufferView body + Just <$> evaluate (Snappy.compress page) other -> error ("writeParquet: unsupported codec " <> show other) let compressedPageSize = maybe uncompressedPageSize BS.length compressedBody headerBytes = @@ -339,19 +328,16 @@ writeDataPage codec numPageRows body columnChunkState = do (mkDataPageHeader numPageRows uncompressedPageSize compressedPageSize) writeByteString columnChunkState.buffer headerBytes case compressedBody of - Nothing -> flushBufferToBuffer body columnChunkState.buffer + Nothing -> stToIO (flushBufferToBuffer body columnChunkState.buffer) Just bytes -> writeByteString columnChunkState.buffer bytes - modifyMutVar' - columnChunkState.uncompressedBufferSize - (+ fromIntegral (BS.length headerBytes + uncompressedPageSize)) + stToIO $ + modifySTRef' + columnChunkState.uncompressedBufferSize + (+ fromIntegral (BS.length headerBytes + uncompressedPageSize)) -flushRowGroup :: - (PrimBase m, MonadIO m) => - ParquetWriteOptions -> - ParquetWriterState m -> - m () +flushRowGroup :: ParquetWriteOptions -> ParquetWriterState -> IO () flushRowGroup options writerState = do - rowNumber <- readMutVar writerState.rowNumberRef + rowNumber <- readIORef writerState.rowNumberRef when (rowNumber > 0) $ do VB.forM_ writerState.columnChunks @@ -359,22 +345,19 @@ flushRowGroup options writerState = do (reversedColumnChunks, totalCompressed, totalUncompressed) <- VB.foldM' ( \(acc, totalCompressedSize, totalUncompressedSize) columnChunkState -> do - offset <- - readMutVar writerState.currentFileOffsetRef - compressedSize <- - bufferResidency columnChunkState.buffer - uncompressedSize <- - readMutVar - columnChunkState.uncompressedBufferSize + offset <- readIORef writerState.currentFileOffsetRef + (compressedSize, uncompressedSize) <- + stToIO $ + (,) + <$> bufferResidency columnChunkState.buffer + <*> readSTRef columnChunkState.uncompressedBufferSize flushBufferToFile writerState.outputFileHandle columnChunkState.buffer - writeMutVar + writeIORef writerState.currentFileOffsetRef (offset + fromIntegral compressedSize) - writeMutVar - columnChunkState.uncompressedBufferSize - 0 + stToIO (writeSTRef columnChunkState.uncompressedBufferSize 0) let columnChunk = mkColumnChunk options.compressionCodec @@ -392,7 +375,7 @@ flushRowGroup options writerState = do ) ([], 0 :: Int64, 0 :: Int64) writerState.columnChunks - modifyMutVar' + modifyIORef' writerState.rowGroupMetadataRef ( mkRowGroup (reverse reversedColumnChunks) @@ -401,12 +384,9 @@ flushRowGroup options writerState = do rowNumber : ) - writeMutVar writerState.rowNumberRef 0 + writeIORef writerState.rowNumberRef 0 -bufferedSize :: - (PrimMonad m) => - VB.Vector (ColumnChunkState m) -> - m Int +bufferedSize :: VB.Vector (ColumnChunkState s) -> ST s Int bufferedSize = VB.foldM' ( \total columnChunkState -> do @@ -421,11 +401,7 @@ bufferedSize = 0 initColumnChunkState :: - (PrimBase m, MonadIO m) => - ParquetWriteOptions -> - T.Text -> - Column -> - m (ColumnChunkState m) + ParquetWriteOptions -> T.Text -> Column -> ST s (ColumnChunkState s) initColumnChunkState options columnName_ column = do encoder_ <- buildEncoder column let nullable_ = hasMissing column @@ -450,7 +426,7 @@ initColumnChunkState options columnName_ column = do -- is likely to hit the page limit, the others are liable to be -- much smaller than the limit. buffer_ <- mallocBuffer bufferSize - uncompressedBufferSize_ <- newMutVar 0 + uncompressedBufferSize_ <- newSTRef 0 pageState_ <- initPageState bufferSize pure ColumnChunkState @@ -463,11 +439,11 @@ initColumnChunkState options columnName_ column = do , pageState = pageState_ } -initPageState :: (PrimMonad m, MonadIO m) => Int -> m (PageState m) +initPageState :: Int -> ST s (PageState s) initPageState bufferSize = do pageBuffer_ <- mallocBuffer bufferSize definitionLevels_ <- newDefLevels - currentRowCount_ <- newMutVar 0 + currentRowCount_ <- newSTRef 0 pure PageState { pageBuffer = pageBuffer_ diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/DefLevels.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/DefLevels.hs index 65efe19b..6dea2138 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/DefLevels.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/DefLevels.hs @@ -8,53 +8,52 @@ module DataFrame.IO.Parquet.Writer.DefLevels ( ) where import Control.Monad (when) -import Control.Monad.IO.Class (MonadIO) -import Control.Monad.Primitive (PrimMonad, PrimState) +import Control.Monad.ST (ST) import Data.Bits (shiftL, shiftR, (.&.), (.|.)) -import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, writeMutVar) +import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef) import Data.Word (Word64) import DataFrame.IO.Utils.RandomAccess (MemoryBuffer, mallocBuffer, writeWord8) data DefLevels s = DefLevels { dlBuf :: !(MemoryBuffer s) - , dlValue :: !(MutVar s Int) - , dlCount :: !(MutVar s Int) + , dlValue :: !(STRef s Int) + , dlCount :: !(STRef s Int) } -newDefLevels :: (PrimMonad m, MonadIO m) => m (DefLevels (PrimState m)) -newDefLevels = DefLevels <$> mallocBuffer 64 <*> newMutVar 0 <*> newMutVar 0 +newDefLevels :: ST s (DefLevels s) +newDefLevels = DefLevels <$> mallocBuffer 64 <*> newSTRef 0 <*> newSTRef 0 -pushDef :: (PrimMonad m) => DefLevels (PrimState m) -> Int -> m () +pushDef :: DefLevels s -> Int -> ST s () pushDef dl value = do - count <- readMutVar dl.dlCount + count <- readSTRef dl.dlCount if count == 0 - then writeMutVar dl.dlValue value >> writeMutVar dl.dlCount 1 + then writeSTRef dl.dlValue value >> writeSTRef dl.dlCount 1 else do - current <- readMutVar dl.dlValue + current <- readSTRef dl.dlValue if current == value - then writeMutVar dl.dlCount (count + 1) + then writeSTRef dl.dlCount (count + 1) else do writeDefRun dl current count - writeMutVar dl.dlValue value - writeMutVar dl.dlCount 1 + writeSTRef dl.dlValue value + writeSTRef dl.dlCount 1 {-# INLINE pushDef #-} -flushDef :: (PrimMonad m) => DefLevels (PrimState m) -> m () +flushDef :: DefLevels s -> ST s () flushDef dl = do - count <- readMutVar dl.dlCount + count <- readSTRef dl.dlCount when (count > 0) $ do - value <- readMutVar dl.dlValue + value <- readSTRef dl.dlValue writeDefRun dl value count - writeMutVar dl.dlCount 0 + writeSTRef dl.dlCount 0 {-# INLINE flushDef #-} -writeDefRun :: (PrimMonad m) => DefLevels (PrimState m) -> Int -> Int -> m () +writeDefRun :: DefLevels s -> Int -> Int -> ST s () writeDefRun dl value count = do writeLeb128 dl.dlBuf (fromIntegral (count `shiftL` 1)) writeWord8 dl.dlBuf (fromIntegral value) {-# INLINE writeDefRun #-} -writeLeb128 :: (PrimMonad m) => MemoryBuffer (PrimState m) -> Word64 -> m () +writeLeb128 :: MemoryBuffer s -> Word64 -> ST s () writeLeb128 buffer value | value < 0x80 = writeWord8 buffer (fromIntegral value) | otherwise = do diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs index 3b5b6afe..c0715da8 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Encoder.hs @@ -4,20 +4,17 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} -{-# OPTIONS_GHC -funfolding-use-threshold=1000 #-} module DataFrame.IO.Parquet.Writer.Encoder ( Encoder (..), buildEncoder, ) where -import Control.Monad.IO.Class (MonadIO, liftIO) -import Control.Monad.Primitive (PrimBase, PrimMonad, PrimState) -import Control.Monad.ST (stToIO) +import Control.Monad.ST (ST) import Data.Bits (shiftL, (.|.)) import Data.Int (Int32, Int64) import Data.Primitive.ByteArray (withMutableByteArrayContents, writeByteArray) -import Data.Primitive.MutVar (newMutVar, readMutVar, writeMutVar) +import Data.STRef (newSTRef, readSTRef, writeSTRef) import qualified Data.Text as T import qualified Data.Text.Array as TA import Data.Text.Internal (Text (Text)) @@ -55,68 +52,76 @@ import GHC.Float (castDoubleToWord64, castFloatToWord32) import Pinch (enum, putField) import Type.Reflection (typeRep) -data Encoder m = Encoder +data Encoder s = Encoder { encType :: !ThriftType , convertedType :: !(Maybe ConvertedType) , logicalType :: !(Maybe LogicalType) - , encodeValue :: - !(MemoryBuffer (PrimState m) -> Int -> Int -> m (Int, Bool)) - , finishValues :: !(MemoryBuffer (PrimState m) -> Int -> m Int) + , encodeValue :: !(MemoryBuffer s -> Int -> Int -> ST s (Int, Bool)) + , finishValues :: !(MemoryBuffer s -> Int -> ST s Int) } -buildEncoder :: (PrimBase m, MonadIO m) => Column -> m (Encoder m) +buildEncoder :: Column -> ST s (Encoder s) buildEncoder col | hasElemType @Int32 col = pure $ - scalarEncoder @Int32 + scalarEncoder (INT32 enum) Nothing Nothing - (\buffer pos v -> write32 buffer pos (fromIntegral v)) - col + ( unboxedColumnWriter @Int32 + col + (\buffer pos v -> write32 buffer pos (fromIntegral v)) + ) | hasElemType @Int64 col = pure $ - scalarEncoder @Int64 + scalarEncoder (INT64 enum) Nothing Nothing - (\buffer pos v -> write64 buffer pos (fromIntegral v)) - col + ( unboxedColumnWriter @Int64 + col + (\buffer pos v -> write64 buffer pos (fromIntegral v)) + ) -- Ints in GHC can be 32 bit or 64 bit integers depending on the -- underlying computers architecture. So we'll do 64bit integers -- to cover all our bases | hasElemType @Int col = pure $ - scalarEncoder @Int + scalarEncoder (INT64 enum) Nothing Nothing - (\buffer pos v -> write64 buffer pos (fromIntegral v)) - col + ( unboxedColumnWriter @Int + col + (\buffer pos v -> write64 buffer pos (fromIntegral v)) + ) | hasElemType @Integer col = pure $ - scalarEncoder @Integer + scalarEncoder (INT64 enum) Nothing Nothing - writeInteger64At - col + (columnWriter @Integer col writeInteger64At) | hasElemType @Float col = pure $ - scalarEncoder @Float + scalarEncoder (FLOAT enum) Nothing Nothing - (\buffer pos v -> write32 buffer pos (castFloatToWord32 v)) - col + ( unboxedColumnWriter @Float + col + (\buffer pos v -> write32 buffer pos (castFloatToWord32 v)) + ) | hasElemType @Double col = pure $ - scalarEncoder @Double + scalarEncoder (DOUBLE enum) Nothing Nothing - (\buffer pos v -> write64 buffer pos (castDoubleToWord64 v)) - col + ( unboxedColumnWriter @Double + col + (\buffer pos v -> write64 buffer pos (castDoubleToWord64 v)) + ) | hasElemType @Bool col = boolEncoder col | hasElemType @T.Text col = pure (textEncoder col) | hasElemType @UTCTime col = pure (timestampEncoder col) @@ -130,36 +135,58 @@ buildEncoder col writeWord64At buffer pos value pure (pos + 8) -{-# SPECIALIZE buildEncoder :: Column -> IO (Encoder IO) #-} +type ColumnWriter s = MemoryBuffer s -> Int -> Int -> ST s (Int, Bool) scalarEncoder :: - forall a m. - (Columnable a, Monad m) => ThriftType -> Maybe ConvertedType -> Maybe LogicalType -> - (MemoryBuffer (PrimState m) -> Int -> a -> m Int) -> - Column -> - Encoder m -scalarEncoder tt conv logical writePrim col = - Encoder tt conv logical (columnWriter @a col writePrim) (\_ pos -> pure pos) + ColumnWriter s -> + Encoder s +scalarEncoder tt conv logical encode = + Encoder tt conv logical encode (\_ pos -> pure pos) + +-- | Encode a column whose element type has no 'VU.Unbox' instance. columnWriter :: - forall a m. - (Columnable a, Monad m) => + forall a s. + (Columnable a) => + Column -> + (MemoryBuffer s -> Int -> a -> ST s Int) -> + ColumnWriter s +columnWriter = columnWriterWith Nothing +{-# INLINE columnWriter #-} + +{- | Encode a column whose element type has a 'VU.Unbox' instance. Using +the caller's instance, rather than the one stored in the column, lets GHC +read the array directly when the element type is known. +-} +unboxedColumnWriter :: + forall a s. + (Columnable a, VU.Unbox a) => Column -> - (MemoryBuffer (PrimState m) -> Int -> a -> m Int) -> - MemoryBuffer (PrimState m) -> - Int -> - Int -> - m (Int, Bool) -columnWriter col writePrim = case col of + (MemoryBuffer s -> Int -> a -> ST s Int) -> + ColumnWriter s +unboxedColumnWriter = columnWriterWith (Just VU.unsafeIndex) +{-# INLINE unboxedColumnWriter #-} + +columnWriterWith :: + forall a s. + (Columnable a) => + Maybe (VU.Vector a -> Int -> a) -> + Column -> + (MemoryBuffer s -> Int -> a -> ST s Int) -> + ColumnWriter s +columnWriterWith unboxedIndex col writePrim = case col of BoxedColumn bitmap (values :: VB.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of Just Refl -> writeFrom bitmap (VB.unsafeIndex values) Nothing -> mismatch UnboxedColumn bitmap (values :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of - Just Refl -> writeFrom bitmap (VU.unsafeIndex values) + Just Refl -> + writeFrom + bitmap + (maybe (VU.unsafeIndex values) ($ values) unboxedIndex) Nothing -> mismatch _ -> mismatch where @@ -171,53 +198,56 @@ columnWriter col writePrim = case col of mismatch = error ("writeParquet: incompatible column representation for " <> columnTypeString col) -{-# INLINE columnWriter #-} +{-# INLINE columnWriterWith #-} isPresent :: Maybe Bitmap -> Int -> Bool isPresent Nothing _ = True isPresent (Just bitmap) row = bitmapTestBit bitmap row {-# INLINE isPresent #-} -boolEncoder :: (PrimMonad m) => Column -> m (Encoder m) +boolEncoder :: Column -> ST s (Encoder s) boolEncoder col = do - bitsRef <- newMutVar (0 :: Word8) - countRef <- newMutVar (0 :: Int) + bitsRef <- newSTRef (0 :: Word8) + countRef <- newSTRef (0 :: Int) let addBit buffer pos value = do - bits <- readMutVar bitsRef - count <- readMutVar countRef + bits <- readSTRef bitsRef + count <- readSTRef countRef let bits' = if value then bits .|. ((1 :: Word8) `shiftL` count) else bits count' = count + 1 if count' == 8 then do - arr <- readMutVar buffer.arrayRef + arr <- readSTRef buffer.arrayRef writeByteArray arr pos bits' - writeMutVar bitsRef 0 - writeMutVar countRef 0 + writeSTRef bitsRef 0 + writeSTRef countRef 0 pure (pos + 1) else do - writeMutVar bitsRef bits' - writeMutVar countRef count' + writeSTRef bitsRef bits' + writeSTRef countRef count' pure pos finish buffer pos = do - count <- readMutVar countRef + count <- readSTRef countRef pos' <- if count > 0 then do - bits <- readMutVar bitsRef - arr <- readMutVar buffer.arrayRef + bits <- readSTRef bitsRef + arr <- readSTRef buffer.arrayRef writeByteArray arr pos bits pure (pos + 1) else pure pos - writeMutVar bitsRef 0 - writeMutVar countRef 0 + writeSTRef bitsRef 0 + writeSTRef countRef 0 pure pos' pure - (Encoder (BOOLEAN enum) Nothing Nothing (columnWriter @Bool col addBit) finish) + ( Encoder + (BOOLEAN enum) + Nothing + Nothing + (unboxedColumnWriter @Bool col addBit) + finish + ) -textEncoder :: - (PrimBase m, MonadIO m) => - Column -> - Encoder m +textEncoder :: Column -> Encoder s textEncoder col = Encoder (BYTE_ARRAY enum) @@ -248,28 +278,18 @@ textEncoder col = pure (pos', True) | otherwise = pure (pos, False) writeTextSlice buffer pos bytes offset count = do - writeMutVar buffer.positionRef pos + writeSTRef buffer.positionRef pos _ <- ensureCapacity buffer (pos + 4 + count) writeWord32At buffer pos (fromIntegral count) - arr <- readMutVar buffer.arrayRef + arr <- readSTRef buffer.arrayRef withMutableByteArrayContents arr $ \ptr -> - liftIO $ - stToIO - ( TA.copyToPointer - bytes - offset - (ptr `plusPtr` (pos + 4)) - count - ) + TA.copyToPointer bytes offset (ptr `plusPtr` (pos + 4)) count pure (pos + 4 + count) mismatch = error ("writeParquet: incompatible text representation for " <> columnTypeString col) -timestampEncoder :: - (PrimMonad m) => - Column -> - Encoder m +timestampEncoder :: Column -> Encoder s timestampEncoder col = Encoder (INT64 enum) diff --git a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Metadata.hs b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Metadata.hs index 6092b9bc..f898489a 100644 --- a/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Metadata.hs +++ b/dataframe-parquet/src/DataFrame/IO/Parquet/Writer/Metadata.hs @@ -1,4 +1,3 @@ -{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} module DataFrame.IO.Parquet.Writer.Metadata ( @@ -11,8 +10,7 @@ module DataFrame.IO.Parquet.Writer.Metadata ( magic, ) where -import Control.Monad.IO.Class (MonadIO) -import Control.Monad.Primitive (PrimBase) +import Control.Monad.ST (stToIO) import qualified Data.ByteString as BS import Data.Int (Int64) import qualified Data.Text as T @@ -140,13 +138,12 @@ mkRowGroup chunks totalCompressed totalUncompressed rgRows = } writeFooter :: - (PrimBase m, MonadIO m) => WritableBinaryHandle -> [SchemaElement] -> Int -> [RowGroup] -> [(T.Text, T.Text)] -> - m () + IO () writeFooter output schemaElements numRows rowGroupMetadata keyValues = do let metadata = FileMetadata @@ -169,9 +166,9 @@ writeFooter output schemaElements numRows rowGroupMetadata keyValues = do , footer_signing_key_metadata = putField Nothing } footer = Pinch.encode Pinch.compactProtocol metadata - buffer <- mallocBuffer (BS.length footer + 8) + buffer <- stToIO (mallocBuffer (BS.length footer + 8)) writeByteString buffer footer - writeWord32LE buffer (fromIntegral (BS.length footer)) + stToIO (writeWord32LE buffer (fromIntegral (BS.length footer))) writeByteString buffer magic flushBufferToFile output buffer diff --git a/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs b/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs index 30d86bae..08b868b2 100644 --- a/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs +++ b/dataframe-parquet/src/DataFrame/IO/Utils/RandomAccess.hs @@ -1,4 +1,5 @@ {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} {-# LANGUAGE OverloadedRecordDot #-} module DataFrame.IO.Utils.RandomAccess ( @@ -29,31 +30,31 @@ module DataFrame.IO.Utils.RandomAccess ( writeDoubleLE, bufferResidency, bufferToByteString, + unsafeBufferView, flushBufferToBuffer, resetPosition, flushBufferToFile, writeByteStringToFile, ) where -import Control.Exception (bracket, bracketOnError, finally) +import Control.Exception (bracket, bracketOnError, finally, throw) import Control.Monad (when) import Control.Monad.IO.Class (MonadIO (..)) -import Control.Monad.Primitive (PrimBase, PrimMonad, PrimState) -import Control.Monad.ST (stToIO) +import Control.Monad.ST (RealWorld, ST, stToIO) import Data.Bits (shiftR) -import qualified Data.ByteString as BS -import Data.ByteString.Internal (ByteString (PS), create) +import Data.ByteString.Internal (ByteString (BS, PS), create) import qualified Data.ByteString.Unsafe as BU import Data.Int (Int64) import Data.Primitive.ByteArray ( - MutableByteArray, + MutableByteArray (..), copyMutableByteArray, getSizeofMutableByteArray, + mutableByteArrayContents, newPinnedByteArray, withMutableByteArrayContents, writeByteArray, ) -import Data.Primitive.MutVar (MutVar, newMutVar, readMutVar, writeMutVar) +import Data.STRef (STRef, newSTRef, readSTRef, writeSTRef) import qualified Data.Text.Array as TA import qualified Data.Vector.Storable as VS import Data.Word (Word32, Word64, Word8) @@ -65,6 +66,8 @@ import DataFrame.IO.Parquet.Seeking ( ) import Foreign (castForeignPtr, castPtr, copyBytes, plusPtr) import GHC.Float (castDoubleToWord64, castFloatToWord32) +import GHC.ForeignPtr (ForeignPtr (..), ForeignPtrContents (PlainPtr)) +import GHC.Ptr (Ptr (..)) import System.Directory (copyPermissions, doesFileExist, removeFile, renameFile) import System.FilePath (takeDirectory) import System.IO ( @@ -201,17 +204,16 @@ withWritableBinaryFile filepath = (hClose . unHandle) data MemoryBuffer s = MemoryBuffer - { arrayRef :: !(MutVar s (MutableByteArray s)) - , positionRef :: !(MutVar s Int) + { arrayRef :: !(STRef s (MutableByteArray s)) + , positionRef :: !(STRef s Int) } -mallocBuffer :: - (PrimMonad m, MonadIO m) => Int -> m (MemoryBuffer (PrimState m)) +mallocBuffer :: Int -> ST s (MemoryBuffer s) mallocBuffer capacity - | capacity < 0 = liftIO $ ioError $ userError "mallocBuffer: negative capacity" + | capacity < 0 = throw (userError "mallocBuffer: negative capacity") | otherwise = do array <- newPinnedByteArray capacity - MemoryBuffer <$> newMutVar array <*> newMutVar 0 + MemoryBuffer <$> newSTRef array <*> newSTRef 0 -- We're using pinned ByteArrays so we must -- not use the grow function brovided by Data.Primitive @@ -228,65 +230,54 @@ mallocBuffer capacity -- is just a matter of adding a new buffer to the array (which we can -- pre-allocate to three elements to begin with and grow it only on the -- off chance that a buffer required more than three grows). -ensureCapacity :: - (PrimMonad m) => - MemoryBuffer (PrimState m) -> Int -> m (MutableByteArray (PrimState m)) +ensureCapacity :: MemoryBuffer s -> Int -> ST s (MutableByteArray s) ensureCapacity buffer needed = do - array <- readMutVar buffer.arrayRef + array <- readSTRef buffer.arrayRef maxSize <- getSizeofMutableByteArray array if needed <= maxSize then pure array else do - position <- readMutVar buffer.positionRef + position <- readSTRef buffer.positionRef grown <- newPinnedByteArray (needed + (needed `div` 2)) copyMutableByteArray grown 0 array 0 position - writeMutVar buffer.arrayRef grown + writeSTRef buffer.arrayRef grown pure grown {-# INLINE ensureCapacity #-} -writeWord8 :: (PrimMonad m) => MemoryBuffer (PrimState m) -> Word8 -> m () +writeWord8 :: MemoryBuffer s -> Word8 -> ST s () writeWord8 buffer b = do - position <- readMutVar buffer.positionRef + position <- readSTRef buffer.positionRef array <- ensureCapacity buffer (position + 1) writeByteArray array position b - writeMutVar buffer.positionRef (position + 1) + writeSTRef buffer.positionRef (position + 1) {-# INLINE writeWord8 #-} -writeByteString :: - (PrimBase m, MonadIO m) => - MemoryBuffer (PrimState m) -> - ByteString -> - m () -writeByteString buffer bs = do - position <- readMutVar buffer.positionRef - let len = BS.length bs - array <- ensureCapacity buffer (position + len) - withMutableByteArrayContents array $ \dst -> - liftIO $ - BU.unsafeUseAsCStringLen bs $ \(source, _) -> do - copyBytes - (dst `plusPtr` position) - (castPtr source) - len - writeMutVar buffer.positionRef (position + len) +-- Reading the ByteString's pointer needs IO, so this stays out of ST. +writeByteString :: MemoryBuffer RealWorld -> ByteString -> IO () +writeByteString buffer bs = + BU.unsafeUseAsCStringLen bs $ \(source, len) -> do + position <- stToIO (readSTRef buffer.positionRef) + array <- stToIO (ensureCapacity buffer (position + len)) + withMutableByteArrayContents array $ \dst -> + copyBytes (dst `plusPtr` position) (castPtr source) len + stToIO (writeSTRef buffer.positionRef (position + len)) {-# INLINE writeByteString #-} -writeWord32LE :: (PrimMonad m) => MemoryBuffer (PrimState m) -> Word32 -> m () +writeWord32LE :: MemoryBuffer s -> Word32 -> ST s () writeWord32LE buffer w = do - position <- readMutVar buffer.positionRef + position <- readSTRef buffer.positionRef writeWord32At buffer position w - writeMutVar buffer.positionRef (position + 4) + writeSTRef buffer.positionRef (position + 4) {-# INLINE writeWord32LE #-} -writeWord64LE :: (PrimMonad m) => MemoryBuffer (PrimState m) -> Word64 -> m () +writeWord64LE :: MemoryBuffer s -> Word64 -> ST s () writeWord64LE buffer w = do - position <- readMutVar buffer.positionRef + position <- readSTRef buffer.positionRef writeWord64At buffer position w - writeMutVar buffer.positionRef (position + 8) + writeSTRef buffer.positionRef (position + 8) {-# INLINE writeWord64LE #-} -writeWord32At :: - (PrimMonad m) => MemoryBuffer (PrimState m) -> Int -> Word32 -> m () +writeWord32At :: MemoryBuffer s -> Int -> Word32 -> ST s () writeWord32At buffer position w = do array <- ensureCapacity buffer (position + 4) writeByteArray array position (fromIntegral w :: Word8) @@ -295,8 +286,7 @@ writeWord32At buffer position w = do writeByteArray array (position + 3) (fromIntegral (w `shiftR` 24) :: Word8) {-# INLINE writeWord32At #-} -writeWord64At :: - (PrimMonad m) => MemoryBuffer (PrimState m) -> Int -> Word64 -> m () +writeWord64At :: MemoryBuffer s -> Int -> Word64 -> ST s () writeWord64At buffer position w = do array <- ensureCapacity buffer (position + 8) writeByteArray array position (fromIntegral w :: Word8) @@ -309,17 +299,14 @@ writeWord64At buffer position w = do writeByteArray array (position + 7) (fromIntegral (w `shiftR` 56) :: Word8) {-# INLINE writeWord64At #-} -writeInteger64 :: - (PrimMonad m, MonadIO m) => MemoryBuffer (PrimState m) -> Integer -> m () +writeInteger64 :: MemoryBuffer s -> Integer -> ST s () writeInteger64 buffer value = do - position <- readMutVar buffer.positionRef + position <- readSTRef buffer.positionRef newPosition <- writeInteger64At buffer position value - writeMutVar buffer.positionRef newPosition + writeSTRef buffer.positionRef newPosition {-# INLINE writeInteger64 #-} -writeInteger64At :: - (PrimMonad m, MonadIO m) => - MemoryBuffer (PrimState m) -> Int -> Integer -> m Int +writeInteger64At :: MemoryBuffer s -> Int -> Integer -> ST s Int writeInteger64At buffer position value | value < toInteger (minBound :: Int64) = outOfRange | value > toInteger (maxBound :: Int64) = outOfRange @@ -328,27 +315,24 @@ writeInteger64At buffer position value pure (position + 8) where outOfRange = - liftIO - (ioError (userError "writeParquet: Integer value is outside the INT64 range")) + throw (userError "writeParquet: Integer value is outside the INT64 range") {-# INLINE writeInteger64At #-} -writeFloatLE :: (PrimMonad m) => MemoryBuffer (PrimState m) -> Float -> m () +writeFloatLE :: MemoryBuffer s -> Float -> ST s () writeFloatLE buffer = writeWord32LE buffer . castFloatToWord32 {-# INLINE writeFloatLE #-} -writeDoubleLE :: (PrimMonad m) => MemoryBuffer (PrimState m) -> Double -> m () +writeDoubleLE :: MemoryBuffer s -> Double -> ST s () writeDoubleLE buffer = writeWord64LE buffer . castDoubleToWord64 {-# INLINE writeDoubleLE #-} -flushBufferToBuffer :: - (PrimMonad m) => - MemoryBuffer (PrimState m) -> MemoryBuffer (PrimState m) -> m () +flushBufferToBuffer :: MemoryBuffer s -> MemoryBuffer s -> ST s () flushBufferToBuffer source destination | source.arrayRef == destination.arrayRef = pure () | otherwise = do - sourceArray <- readMutVar source.arrayRef - sourcePosition <- readMutVar source.positionRef - destinationPosition <- readMutVar destination.positionRef + sourceArray <- readSTRef source.arrayRef + sourcePosition <- readSTRef source.positionRef + destinationPosition <- readSTRef destination.positionRef destinationArray <- ensureCapacity destination (destinationPosition + sourcePosition) copyMutableByteArray @@ -357,28 +341,35 @@ flushBufferToBuffer source destination sourceArray 0 sourcePosition - writeMutVar destination.positionRef (destinationPosition + sourcePosition) - writeMutVar source.positionRef 0 + writeSTRef destination.positionRef (destinationPosition + sourcePosition) + writeSTRef source.positionRef 0 {-# INLINE flushBufferToBuffer #-} -bufferToByteString :: - (PrimBase m, MonadIO m) => - MemoryBuffer (PrimState m) -> - m ByteString +bufferToByteString :: MemoryBuffer RealWorld -> IO ByteString bufferToByteString buffer = do - array <- readMutVar buffer.arrayRef - position <- readMutVar buffer.positionRef - withMutableByteArrayContents array $ \src -> - liftIO $ - create position $ \dst -> - copyBytes dst (castPtr src) position - -bufferResidency :: (PrimMonad m) => MemoryBuffer (PrimState m) -> m Int -bufferResidency buffer = readMutVar buffer.positionRef + array <- stToIO (readSTRef buffer.arrayRef) + position <- stToIO (readSTRef buffer.positionRef) + create position $ \dst -> + withMutableByteArrayContents array $ \src -> + copyBytes dst (castPtr src) position + +{- | The buffer's contents as a ByteString, without copying. The result +shares memory with the buffer, so it is only valid until the buffer is +next written to or reset. +-} +unsafeBufferView :: MemoryBuffer RealWorld -> IO ByteString +unsafeBufferView buffer = do + array@(MutableByteArray array#) <- stToIO (readSTRef buffer.arrayRef) + position <- stToIO (readSTRef buffer.positionRef) + let Ptr address# = mutableByteArrayContents array + pure (BS (ForeignPtr address# (PlainPtr array#)) position) + +bufferResidency :: MemoryBuffer s -> ST s Int +bufferResidency buffer = readSTRef buffer.positionRef {-# INLINE bufferResidency #-} -resetPosition :: (PrimMonad m) => MemoryBuffer (PrimState m) -> m () -resetPosition buffer = writeMutVar buffer.positionRef 0 +resetPosition :: MemoryBuffer s -> ST s () +resetPosition buffer = writeSTRef buffer.positionRef 0 {-# INLINE resetPosition #-} -- I tested write speeds by doing (on Apple Silicon) @@ -402,13 +393,11 @@ resetPosition buffer = writeMutVar buffer.positionRef 0 -- So when writing to a file to minimize syscall overhead while -- trying not to create dirty pages in the kernel page cache, we'll -- be flushing in 256 KiB chunks. -flushBufferToFile :: - (PrimBase m, MonadIO m) => - WritableBinaryHandle -> MemoryBuffer (PrimState m) -> m () +flushBufferToFile :: WritableBinaryHandle -> MemoryBuffer RealWorld -> IO () flushBufferToFile (WritableBinaryHandle h) buffer = do - array <- readMutVar buffer.arrayRef - position <- readMutVar buffer.positionRef - withMutableByteArrayContents array $ \ptr -> liftIO $ do + array <- stToIO (readSTRef buffer.arrayRef) + position <- stToIO (readSTRef buffer.positionRef) + withMutableByteArrayContents array $ \ptr -> do let chunkSize = 262144 go offset | offset >= position = pure () @@ -417,7 +406,7 @@ flushBufferToFile (WritableBinaryHandle h) buffer = do hPutBuf h (ptr `plusPtr` offset) n go (offset + n) go 0 - writeMutVar buffer.positionRef 0 + stToIO (resetPosition buffer) writeByteStringToFile :: WritableBinaryHandle -> ByteString -> IO () writeByteStringToFile (WritableBinaryHandle h) bs = @@ -431,23 +420,13 @@ writeByteStringToFile (WritableBinaryHandle h) bs = go (offset + n) go 0 -appendTextArraySlice :: - (PrimBase m, MonadIO m) => - MemoryBuffer (PrimState m) -> TA.Array -> Int -> Int -> m () +appendTextArraySlice :: MemoryBuffer s -> TA.Array -> Int -> Int -> ST s () appendTextArraySlice buffer source offset count - | count < 0 = - liftIO $ ioError $ userError "appendTextArraySlice: negative length" + | count < 0 = throw (userError "appendTextArraySlice: negative length") | otherwise = do - position <- readMutVar buffer.positionRef + position <- readSTRef buffer.positionRef array <- ensureCapacity buffer (position + count) withMutableByteArrayContents array $ \destination -> - liftIO $ - stToIO - ( TA.copyToPointer - source - offset - (destination `plusPtr` position) - count - ) - writeMutVar buffer.positionRef (position + count) + TA.copyToPointer source offset (destination `plusPtr` position) count + writeSTRef buffer.positionRef (position + count) {-# INLINE appendTextArraySlice #-} diff --git a/dataframe-parquet/tests/Main.hs b/dataframe-parquet/tests/Main.hs index fcbfe425..73fffc63 100644 --- a/dataframe-parquet/tests/Main.hs +++ b/dataframe-parquet/tests/Main.hs @@ -12,6 +12,7 @@ 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 @@ -34,14 +35,14 @@ import System.Directory (listDirectory) directWrites :: Test directWrites = TestCase $ do - buffer <- mallocBuffer 1 - writeWord8 buffer 0xaa - writeWord32LE buffer 0x78563412 - writeWord64LE buffer 0x0807060504030201 - writeFloatLE buffer 1 - writeDoubleLE buffer 1 + 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 <- bufferResidency buffer + residency <- stToIO (bufferResidency buffer) bytes <- bufferToByteString buffer assertEqual "direct write residency" 27 residency assertEqual @@ -80,20 +81,20 @@ directWrites = TestCase $ do directBufferFlush :: Test directBufferFlush = TestCase $ do - source <- mallocBuffer 0 - destination <- mallocBuffer 0 + source <- stToIO (mallocBuffer 0) + destination <- stToIO (mallocBuffer 0) writeByteString destination (BS.pack [1, 2]) writeByteString source (BS.pack [3, 4, 5]) - flushBufferToBuffer source destination - sourceResidency <- bufferResidency source + 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 - flushBufferToBuffer destination destination + stToIO (flushBufferToBuffer destination destination) selfFlushedBytes <- bufferToByteString destination assertEqual "self flush is a no-op" destinationBytes selfFlushedBytes - resetPosition destination - destinationResidency <- bufferResidency destination + stToIO (resetPosition destination) + destinationResidency <- stToIO (bufferResidency destination) assertEqual "reset position" 0 destinationResidency directFileFlush :: Test @@ -101,11 +102,11 @@ directFileFlush = TestCase $ withSystemTempDirectory "dfpq-buffer" $ \dir -> do let outPath = dir "out.bin" payload = BS.pack (take 300000 (cycle [0 .. 255])) - buffer <- mallocBuffer 1 + buffer <- stToIO (mallocBuffer 1) writeByteString buffer payload withWritableBinaryFile outPath $ \output -> flushBufferToFile output buffer - residency <- bufferResidency buffer + residency <- stToIO (bufferResidency buffer) contents <- BS.readFile outPath assertEqual "source cleared after file flush" 0 residency assertEqual "large payload round-trips" payload contents