如何在Haskell中将大型数据块解析到内存中？

Question

仔细思考后，这个问题可以简化为更加精炼的内容。我正在寻找一个Haskell数据结构，该数据结构具有以下特点：

• 看起来像一个列表
• 拥有O(1)查找
• 具有O(1)元素替换或O(1)元素追加（或者是前置...如果那是情况的话，我可以反转我的索引查找）。我可以随时使用其中之一编写我的后续算法。
• 具有非常少的内存开销

---

我正在尝试构建一个图像文件解析器。文件格式是基本的8位颜色ppm文件，虽然我打算支持16位颜色文件以及PNG和JPEG文件。现有的Netpbm库，尽管有很多未装箱的注释，但在尝试加载我处理的文件时实际上会消耗所有可用内存：

3-10张照片，最小的大小为45MB，最大的大小为110MB。

现在，我无法理解Netpbm代码中的优化，因此我决定自己试试。它是一个简单的文件格式...

我已经开始决定，无论文件格式如何，我都要将最终图像以未压缩的格式存储在其中：

import Data.Vector.Unboxed (Vector)
data PixelMap = RGB8 {
      width :: Int
    , height :: Int
    , redChannel :: Vector Word8
    , greenChannel :: Vector Word8
    , blueChannel :: Vector Word8
    }

然后我编写了一个解析器，它可以像下面这样处理三个向量：

import Data.Attoparsec.ByteString
data Progress = Progress {
      addr      :: Int
    , size      :: Int
    , redC      :: Vector Word8
    , greenC    :: Vector Word8
    , blueC     :: Vector Word8
    }

parseColorBinary :: Progress -> Parser Progress
parseColorBinary progress@Progress{..}
    | addr == size = return progress
    | addr < size = do
        !redV <- anyWord8
        !greenV <- anyWord8
        !blueV <- anyWord8
        parseColorBinary progress { addr    = addr + 1
                                  , redC    = redC V.// [(addr, redV)]
                                  , greenC  = greenC V.// [(addr, greenV)]
                                  , blueC   = blueC V.// [(addr, blueV)] }

在解析器的末尾，我按如下方式构建RGB8：

Progress{..} <- parseColorBinary $ ...
return $ RGB8 width height redC greenC blueC

以这种方式编写的程序，加载其中一个45MB的图像将会消耗5GB或更多的内存。如果我改变Progress的定义，使得redC、greenC和blueC都是!(Vector Word8)，那么程序仍然在合理的内存范围内，但加载单个文件所需的时间太长了，以至于我没有让它完成。最后，如果我用标准列表替换这里的向量，我的内存使用量会飙升到每个文件5GB（我假设...实际上我在达到这个值之前就已经用完了空间），而且加载时间大约为6秒。Ubuntu的预览应用程序一旦启动，几乎瞬间加载并呈现文件。

基于这样的理论，即每次调用V.//实际上都会完全复制向量，我尝试切换到Data.Vector.Unboxed.Mutable，但是...我甚至无法使其类型检查通过。文档不存在，理解数据类型正在做什么需要与多个其他库进行斗争。而且我甚至不知道它是否能解决问题，所以我非常不愿意尝试。

根本问题实际上很简单：如何快速地、不使用过多的内存读取、保留和潜在地操作一个非常大的数据结构？我找到的所有示例都是关于生成临时巨大的数据，然后尽快将其丢弃。

原则上，我希望最终的表示是不可变的，但如果必须使用可变结构才能达到这个目的，我也不太在意。

完整的代码（BSD3许可证）可以在bitbucket上找到：https://bitbucket.org/savannidgerinel/photo-tools。performance分支包含一个严格版本的解析器，可以通过在Codec.Image.Netpbm的Progress数据结构中进行快速更改来使其非严格化。

要运行性能测试：

ulimit -Sv 6000000 -- set a ulimit of 6GB, or change to whatever makes sense for you
cabal build
dist/build/perf-test/perf-test +RTS -p -sstderr

- Savanni D'Gerinel

---

你有没有考虑使用 mmap？如果在图片上执行的操作不是太I/O密集型，那么使用它可能是值得的。 - Danny Navarro

另一个选择是使用一些专门针对大数组的库，比如repa或者accelerate。它们都是为高性能而编写的，因此应该具有许多内存效率优化。 - Danny Navarro

你的45MB图像的像素尺寸是多少？ - Mau

如果你想要帮助让可变向量的类型检查通过，请展示给我们代码。如果能将代码最小化到一个完整的最小代码块（我可以将其放入文件并通过 GHC 运行），并且仍然存在问题，那么会有额外的加分。 - Daniel Wagner

@Mau，不是很大。4767 * 3195 * 3只有45MB。我相信额外的内存，在某些情况下是修改纯数据结构时所需的开销，因此需要一遍又一遍地调用构造函数的所有实例。内存复制的垃圾回收行为非常惊人。 - Savanni D'Gerinel

Answer 1

我最初认为只是简单地读取整个字节串，然后将内容解压缩成未装箱向量就足够了。事实上，即使没有神秘的空间泄漏，您发布的解析代码也相当糟糕：您在输入的每个字节上都复制了所有三个向量的全部内容！谈论二次复杂度。

因此，我编写了以下代码：

chunksOf3 :: [a] -> [(a, a, a)]
chunksOf3 (a:b:c:xs) = (a, b, c) : chunksOf3 xs
chunksOf3 _          = []

parseRGB :: Int -> Atto.Parser (Vector Word8, Vector Word8, Vector Word8)
parseRGB size = do
    input <- Atto.take (size * 3)
    let (rs, gs, bs) = unzip3 $ chunksOf3 $ B.unpack input
    return (V.fromList rs, V.fromList gs, V.fromList bs)

然后我用一个45 Mb的随机字节文件进行了测试。我承认，我很惊讶这段代码导致了数千兆字节的内存使用。我很好奇到底哪里出现了空间泄漏。

可变向量运作良好。以下代码使用133 Mb RAM，Criterion对其进行基准测试，包括60 ms的文件读取。我在注释中添加了一些解释。在SO和其他地方也有关于ST单子和可变向量的丰富资料（虽然我同意库文档对初学者不友好）。

import Data.Vector.Unboxed (Vector)
import Data.ByteString (ByteString)

import qualified Data.Vector.Unboxed as V
import qualified Data.ByteString as B
import qualified Data.Vector.Unboxed.Mutable as MV

import Control.Monad.ST.Strict 
import Data.Word
import Control.Monad
import Control.DeepSeq

-- benchmarking stuff
import Criterion.Main (defaultMainWith, bench, whnfIO)
import Criterion.Config (defaultConfig, Config(..), ljust)

-- This is just the part that parses the three vectors for the colors.
-- Of course, you can embed this into an Attoparsec computation by taking 
-- the current input, feeding it to parseRGB, or you can just take the right 
-- sized chunk in the parser and omit the "Maybe" test from the code below. 
parseRGB :: Int -> ByteString -> Maybe (Vector Word8, Vector Word8, Vector Word8)
parseRGB size input 
    | 3* size > B.length input = Nothing
    | otherwise = Just $ runST $ do

        -- We are allocating three mutable vectors of size "size"
        -- This is usually a bit of pain for new users, because we have to
        -- specify the correct type somewhere, and it's not an exactly simple type.
        -- In the ST monad there is always an "s" type parameter that labels the
        -- state of the action. A type of "ST s something" is a bit similar to
        -- "IO something", except that the inner type often also contains "s" as
        -- parameter. The purpose of that "s" is to statically disallow mutable
        -- variables from escaping the ST action. 
        [r, g, b] <- replicateM 3 $ MV.new size :: ST s [MV.MVector s Word8]

        -- forM_ = flip mapM_
        -- In ST code forM_ is a nicer looking approximation of the usual
        -- imperative loop. 
        forM_ [0..size - 1] $ \i -> do
            let i' = 3 * i
            MV.unsafeWrite r i (B.index input $ i'    )
            MV.unsafeWrite g i (B.index input $ i' + 1)
            MV.unsafeWrite b i (B.index input $ i' + 2)

        -- freeze converts a mutable vector living in the ST monad into 
        -- a regular vector, which can be then returned from the action
        -- since its type no longer depends on that pesky "s".
        -- unsafeFreeze does the conversion in place without copying.
        -- This implies that the original mutable vector should not be used after
        -- unsafeFreezing. 
        [r, g, b] <- mapM V.unsafeFreeze [r, g, b]
        return (r, g, b)

-- I prepared a file with 3 * 15 million random bytes.
inputSize = 15000000
benchConf = defaultConfig {cfgSamples = ljust 10}

main = do
    defaultMainWith benchConf (return ()) $ [
        bench "parseRGB test" $ whnfIO $ do 
            input <- B.readFile "randomInp.dat" 
            force (parseRGB inputSize input) `seq` putStrLn "done"
        ]

- András Kovács

---

我最初使用了Vector，因为我认为结构共享等功能可以防止复制更新操作。但从技术上讲，这是没有意义的。普通列表所带来的所有内存开销可能都是必要的，以允许结构共享。 - Savanni D'Gerinel

似乎我学会了多少库都无关紧要。每次我进入任何新的库时，我总是感觉像一个 Haskell 新手。 - Savanni D'Gerinel

最后，为什么你使用了unsafe操作而不是普通的操作？ - Savanni D'Gerinel

@SavanniD'Gerinel 这只是因为我已经通过初始化建立了边界的正确性，所以我可以跳过边界检查。在大多数编程语言中，我认为不安全索引的广泛使用非常糟糕，但在 Haskell ST 代码中，它相当罕见，而且当我们真正关心速度时才会使用它，所以在这些情况下我可能会尽情发挥 :) - András Kovács

@SavanniD'Gerinel，注意unsafeFreeze这一点更有道理，因为除了在那个时候返回它们之外，我们对向量几乎没有做任何操作。 - András Kovács

Answer 2

这是一个直接从磁盘解析文件而不加载任何中间文件到内存的版本：

import Control.Applicative
import Control.Monad (void)
import Data.Attoparsec.ByteString (anyWord8)
import Data.Attoparsec.ByteString.Char8 (decimal)
import qualified Data.Attoparsec.ByteString as Attoparsec
import Data.ByteString (ByteString)
import Data.Vector.Unboxed (Vector)
import Data.Word (Word8)
import Control.Foldl (FoldM(..), impurely, vector, premapM) -- Uses `foldl-1.0.3`
import qualified Pipes.ByteString
import Pipes.Parse
import Pipes.Attoparsec (parse, parsed)
import qualified System.IO as IO

data PixelMap = PixelMap {
      width :: Int
    , height :: Int
    , redChannel :: Vector Word8
    , greenChannel :: Vector Word8
    , blueChannel :: Vector Word8
    } deriving (Show)

-- Fold three vectors simultaneously, ensuring strictness and efficiency
rgbVectors
    :: FoldM IO (Word8, Word8, Word8) (Vector Word8, Vector Word8, Vector Word8)
rgbVectors =
    (,,) <$> premapM _1 vector <*> premapM _2 vector <*> premapM _3 vector
  where
    _1 (a, b, c) = a
    _2 (a, b, c) = b
    _3 (a, b, c) = c

triples
    :: Monad m
    => Producer ByteString m r
    -> Producer (Word8, Word8, Word8) m ()
triples p = void $ parsed ((,,) <$> anyWord8 <*> anyWord8 <*> anyWord8) p

-- I will probably ask Renzo to simplify the error handling for `parse`
-- This is a helper function to just return `Nothing`
parse'
    :: Monad m
    => Attoparsec.Parser r -> Parser ByteString m (Maybe r)
parse' parser = do
    x <- parse parser
    return $ case x of
        Just (Right r) -> Just r
        _              -> Nothing

parsePixelMap :: Producer ByteString IO r -> IO (Maybe PixelMap)
parsePixelMap p = do
    let parseWH = do
            mw <- parse' decimal
            mh <- parse' decimal
            return ((,) <$> mw <*> mh)
    (x, p') <- runStateT parseWH p
    case x of
        Nothing     -> return Nothing
        Just (w, h) -> do
            let size = w * h
                parser = impurely foldAllM rgbVectors
                source = triples (p' >-> Pipes.ByteString.take size)
            (rs, gs, bs) <- evalStateT parser source
            return $ Just (PixelMap w h rs gs bs)

main = IO.withFile "image.ppm" IO.ReadMode $ \handle -> do
    pixelMap <- parsePixelMap (Pipes.ByteString.fromHandle handle)
    print pixelMap

我测试了一个大小为50 MB的文件，在没有标题逻辑的情况下运行速度与原来差不多。

- Gabriella Gonzalez