Как мне разобрать большой блок данных в памяти в Haskell?

Question

Если подумать, весь этот вопрос можно свести к чему-то гораздо более лаконичному. Я ищу структуру данных Haskell, которая

• выглядит как список
• имеет O(1) поиск
• имеет либо замену элемента O(1), либо добавление элемента O(1) (или prepend... Я мог бы изменить поиск в индексе, если бы это было так). Я всегда могу написать свои более поздние алгоритмы с учетом того или другого.
• имеет очень мало памяти

---

Я пытаюсь создать анализатор файлов изображений. Формат файла - это ваш базовый 8-битный цветной ppm-файл, хотя я намерен поддерживать 16-битные цветные файлы, а также файлы PNG и JPEG. Существующая библиотека Netpbm, несмотря на множество аннотаций для распаковки, фактически использует всю доступную память при попытке загрузить файлы, с которыми я работаю:

3-10 фотографий, самая маленькая из которых составляет 45 МБ, а самая большая - 110 МБ.

Теперь я не могу понять, какие оптимизации вносятся в код 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

Написанная так, программа, загружающая одно из этих изображений размером 45 МБ, будет занимать 5 ГБ или больше памяти. Если я изменю определение Progress чтобы redC, greenC, а также blueC являются все !(Vector Word8), затем программа остается в разумных пределах памяти, но загружается один файл так долго, что я не позволил ему фактически завершить работу. Наконец, если я заменим векторы здесь стандартными списками, мое использование памяти увеличится до 5 ГБ на файл (я полагаю... у меня действительно не хватает места, прежде чем я его ударил), и время загрузки составляет порядка 6 секунд, Предварительно запущенное приложение Ubuntu загружает и отображает файл практически мгновенно.

На теории, что каждый вызов V.// фактически полностью копирует вектор каждый раз, я попытался переключиться на Data.Vector.Unboxed.Mutable, но... я даже не могу это проверить. Документации не существует, и понимание того, что делают типы данных, потребует борьбы и с несколькими другими библиотеками. И я даже не знаю, решит ли это проблемы, поэтому я очень не хочу даже пытаться.

Фундаментальная проблема на самом деле довольно проста:

Как быстро, без использования неприличного объема памяти, читать, сохранять и потенциально манипулировать очень большой структурой данных? Все примеры, которые я нашел, связаны с созданием временно огромных данных и последующим их удалением как можно скорее.

В принципе, я хочу, чтобы окончательное представление было неизменным, но мне не слишком важно, нужно ли мне использовать изменяемую структуру, чтобы туда попасть.

---

Просто для полноты, полный код (BSD3-лицензированный) находится на bitbucket в https://bitbucket.org/savannidgerinel/photo-tools. performance ветвь содержит строгую версию парсера, которая может быть сделана нестрогой с быстрым изменением в Progress структура данных Codec.Image.Netpbm,

Чтобы запустить тест производительности

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

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 МБ файлом случайных байтов. Признаюсь, я был удивлен, что этот код привел к гигабайтам использования оперативной памяти. Мне любопытно, где именно место утечки.

Изменчивые векторы работают хорошо, хотя. В следующем коде используется 133 МБ ОЗУ, а критерий Criterion включает чтение файла до 60 мс. Я включил некоторые пояснения в комментарии. Существует также достаточно материала о монаде ST и изменяемых векторах в SO и в других местах (я согласен с тем, что документация к библиотеке неприветлива для начинающих).

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"
        ]

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 МБ, и он занимает примерно столько же места.