module DataFrame.IO.Parquet.Decompress where

import qualified Codec.Compression.GZip as GZip
import qualified Codec.Compression.Zstd.Base as Zstd
import qualified Data.ByteString as BS
import qualified Data.ByteString as LB
import Data.ByteString.Internal (createAndTrim, toForeignPtr)
import DataFrame.IO.Parquet.Thrift (CompressionCodec (..))
import Foreign.ForeignPtr (withForeignPtr)
import Foreign.Ptr (plusPtr)
import qualified Snappy

decompressData :: Int -> CompressionCodec -> BS.ByteString -> IO BS.ByteString
decompressData :: Int -> CompressionCodec -> ByteString -> IO ByteString
decompressData Int
uncompressedSize CompressionCodec
codec ByteString
compressed = case CompressionCodec
codec of
    (ZSTD Enumeration 6
_) -> Int -> (Ptr Word8 -> IO Int) -> IO ByteString
createAndTrim Int
uncompressedSize ((Ptr Word8 -> IO Int) -> IO ByteString)
-> (Ptr Word8 -> IO Int) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
dstPtr ->
        let (ForeignPtr Word8
srcFP, Int
offset, Int
compressedSize) = ByteString -> (ForeignPtr Word8, Int, Int)
toForeignPtr ByteString
compressed
         in ForeignPtr Word8 -> (Ptr Word8 -> IO Int) -> IO Int
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
srcFP ((Ptr Word8 -> IO Int) -> IO Int)
-> (Ptr Word8 -> IO Int) -> IO Int
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
srcPtr -> do
                Either [Char] Int
result <-
                    Ptr Word8 -> Int -> Ptr Any -> Int -> IO (Either [Char] Int)
forall dst src.
Ptr dst -> Int -> Ptr src -> Int -> IO (Either [Char] Int)
Zstd.decompress
                        Ptr Word8
dstPtr
                        Int
uncompressedSize
                        (Ptr Word8
srcPtr Ptr Word8 -> Int -> Ptr Any
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
offset)
                        Int
compressedSize
                case Either [Char] Int
result of
                    Left [Char]
e -> [Char] -> IO Int
forall a. HasCallStack => [Char] -> a
error ([Char] -> IO Int) -> [Char] -> IO Int
forall a b. (a -> b) -> a -> b
$ [Char]
"ZSTD error: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> [Char]
e
                    Right Int
actualSize -> Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
actualSize
    (SNAPPY Enumeration 1
_) -> case ByteString -> Either DecodeError ByteString
Snappy.decompress ByteString
compressed of
        Left DecodeError
e -> [Char] -> IO ByteString
forall a. HasCallStack => [Char] -> a
error (DecodeError -> [Char]
forall a. Show a => a -> [Char]
show DecodeError
e)
        Right ByteString
res -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
res
    (UNCOMPRESSED Enumeration 0
_) -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
compressed
    (GZIP Enumeration 2
_) -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (LazyByteString -> ByteString
LB.toStrict (LazyByteString -> LazyByteString
GZip.decompress (ByteString -> LazyByteString
BS.fromStrict ByteString
compressed)))
    CompressionCodec
other -> [Char] -> IO ByteString
forall a. HasCallStack => [Char] -> a
error ([Char]
"Unsupported compression type: " [Char] -> [Char] -> [Char]
forall a. Semigroup a => a -> a -> a
<> CompressionCodec -> [Char]
forall a. Show a => a -> [Char]
show CompressionCodec
other)