biohazard-0.6.1: src/Data/Avro.hs
{-# LANGUAGE OverloadedStrings, FlexibleInstances, TemplateHaskell #-}
{-# LANGUAGE RecordWildCards, BangPatterns, FlexibleContexts #-}
module Data.Avro where
import Bio.Iteratee
import Control.Applicative
import Control.Monad
import Control.Monad.ST ( runST, ST )
import Data.Aeson hiding ((.=))
import Data.Array.MArray
import Data.Array.ST ( STUArray )
import Data.Array.Unsafe ( castSTUArray )
import Data.Binary.Get
import Data.Bits
import Data.Binary.Builder
import Data.Foldable ( foldMap )
import Data.Int ( Int64 )
import Data.Maybe
import Data.Monoid
import Data.Scientific
import Data.Text.Encoding
import Data.Word ( Word32, Word64 )
import Foreign.Storable ( Storable, sizeOf )
import Language.Haskell.TH
import System.Random
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.HashMap.Strict as H
import qualified Data.ListLike as LL
import qualified Data.Text as T
import qualified Data.Vector as V
import qualified Data.Vector.Unboxed as U
-- ^ Support for Avro.
-- Current status is that we can generate schemas for certain Haskell
-- values, serialize to binary and JSON representations, and write
-- Container files using the null codec. The C implementation likes
-- some, but not all of these containers; it's unclear if that's the
-- fault of the C implementation, though.
--
-- Meanwhile, serialization works for nested sums-of-products, as long as the
-- product uses record syntax and the top level is a plain record.
-- The obvious primitives are supported.
(.=) :: ToJSON a => String -> a -> (T.Text, Value)
k .= v = (T.pack k, toJSON v)
string :: String -> Value
string = String . T.pack
-- | This is the class of types we can embed into the Avro
-- infrastructure. Right now, we can derive a schema, encode to
-- the Avro binary format, and encode to the Avro JSON encoding.
class Avro a where
-- | Produces the schema for this type. Schemas are represented as
-- JSON values. The monad is used to keep a table of already
-- defined types, so the schema can refer to them by name. (The
-- concrete argument serves to specify the type, it is not actually
-- used.)
toSchema :: a -> MkSchema Value
-- | Serializes a value to the binary representation. The schema is
-- implied, serialization to related schemas is not supported.
toBin :: a -> Builder
-- | Deserializzes a value from binary representation. Right now,
-- no attempt at schema matching is done, the schema must match the
-- expected one exactly.
fromBin :: Get a
-- | Serializes a value to the JSON representation. Note that even
-- the JSON format needs a schema for successful deserialization,
-- and here we support only the one implied schema.
toAvron :: a -> Value
newtype MkSchema a = MkSchema
{ mkSchema :: (a -> H.HashMap T.Text Value -> Value) -> H.HashMap T.Text Value -> Value }
instance Functor MkSchema where fmap f m = MkSchema (\k -> mkSchema m (k . f))
instance Applicative MkSchema where pure a = MkSchema (\k -> k a)
u <*> v = MkSchema (\k -> mkSchema u (\a -> mkSchema v (k . a)))
instance Monad MkSchema where return a = MkSchema (\k -> k a)
a >>= m = MkSchema (\k -> mkSchema a (\a' -> mkSchema (m a') k))
memoObject :: String -> [(T.Text,Value)] -> MkSchema Value
memoObject nm ps = MkSchema $ \k h ->
let nm' = T.pack nm
obj = object $ ("name" .= nm) : ps
in case H.lookup nm' h of
Nothing -> k obj $! H.insert nm' obj h
Just obj' | obj == obj' -> k (String nm') h
| otherwise -> error $ "same type name, different schema: " ++ nm
runMkSchema :: MkSchema Value -> Value
runMkSchema x = mkSchema x postproc H.empty
where
-- Objects are fine as is.
postproc (Object o) _ = Object o
-- Top level can't be a string, can it? Need to wrap into the long form.
postproc (String tp) _ = object [ "type" .= String tp ]
-- Top level Array should be fine, too.
postproc (Array a) _ = Array a
-- reject anything else
postproc v _ = error $ "Not allowed as toplevel schema: " ++ show v
-- instances for primitive types
-- | The Avro \"null\" type is represented as the empty tuple.
instance Avro () where
toSchema _ = return $ String "null"
toBin () = mempty
fromBin = return ()
toAvron () = Null
instance Avro Bool where
toSchema _ = return $ String "boolean"
toBin = singleton . fromIntegral . fromEnum
fromBin = toEnum . fromIntegral <$> getWord8
toAvron = Bool
instance Avro Int where
toSchema _ = return $ String "long"
toBin = encodeIntBase128
fromBin = decodeIntBase128
toAvron = Number . fromIntegral
instance Avro Int64 where
toSchema _ = return $ String "long"
toBin = encodeIntBase128
fromBin = decodeIntBase128
toAvron = Number . fromIntegral
instance Avro Float where
toSchema _ = return $ String "float"
toBin = putWord32le . floatToWord
fromBin = wordToFloat <$> getWord32le
toAvron = Number . fromFloatDigits
instance Avro Double where
toSchema _ = return $ String "double"
toBin = putWord64le . doubleToWord
fromBin = wordToDouble <$> getWord64le
toAvron = Number . fromFloatDigits
instance Avro B.ByteString where
toSchema _ = return $ String "bytes"
toBin s = encodeIntBase128 (B.length s) <> fromByteString s
fromBin = decodeIntBase128 >>= getByteString
toAvron = String . decodeLatin1
instance Avro T.Text where
toSchema _ = return $ String "string"
toBin = toBin . encodeUtf8
fromBin = decodeUtf8 <$> fromBin
toAvron = String
-- Integer<->Float conversions, stolen from cereal.
{-# INLINE wordToFloat #-}
wordToFloat :: Word32 -> Float
wordToFloat x = runST (cast x)
{-# INLINE wordToDouble #-}
wordToDouble :: Word64 -> Double
wordToDouble x = runST (cast x)
{-# INLINE floatToWord #-}
floatToWord :: Float -> Word32
floatToWord x = runST (cast x)
{-# INLINE doubleToWord #-}
doubleToWord :: Double -> Word64
doubleToWord x = runST (cast x)
{-# INLINE cast #-}
cast :: ( MArray (STUArray s) b (ST s), MArray (STUArray s) a (ST s) ) => a -> ST s b
cast x = (newArray (0 :: Int, 0) x >>= castSTUArray >>= flip readArray 0)
-- | Implements Zig-Zag-Coding like in Protocol Buffers and Avro.
zig :: (Storable a, Bits a) => a -> a
zig x = (x `shiftL` 1) `xor` (x `shiftR` (8 * sizeOf x -1))
-- | Reverses Zig-Zag-Coding like in Protocol Buffers and Avro.
zag :: (Storable a, Bits a, Num a) => a -> a
zag x = negate (x .&. 1) `xor` ((x .&. complement 1) `rotateR` 1)
-- | Encodes a word of any size using a variable length "base 128"
-- encoding.
encodeWordBase128 :: (Integral a, Bits a) => a -> Builder
encodeWordBase128 x | x' == 0 = singleton (fromIntegral (x .&. 0x7f))
| otherwise = singleton (fromIntegral (x .&. 0x7f .|. 0x80))
<> encodeWordBase128 x'
where x' = x `shiftR` 7
decodeWordBase128 :: (Integral a, Bits a) => Get a
decodeWordBase128 = go 0 0
where
go acc sc = do x <- getWord8
let !acc' = acc .|. fromIntegral x `shiftL` sc
if x .&. 0x80 == 0
then return acc'
else go acc' (sc+7)
-- | Encodes an int of any size by combining the zig-zag coding with the
-- base 128 encoding.
encodeIntBase128 :: (Integral a, Bits a, Storable a) => a -> Builder
encodeIntBase128 = encodeWordBase128 . zig
-- | Decodes an int of any size by combining the zig-zag decoding with
-- the base 128 decoding.
decodeIntBase128 :: (Integral a, Bits a, Storable a) => Get a
decodeIntBase128 = zag <$> decodeWordBase128
zigInt :: Int -> Builder
zigInt = encodeIntBase128
zagInt :: Get Int
zagInt = decodeWordBase128
-- Complex Types
-- | A list becomes an Avro array
-- The chunked encoding for lists may come in handy. How to select the
-- chunk size is not obvious, though.
instance Avro a => Avro [a] where
toSchema as = do sa <- toSchema (head as)
return $ object [ "type" .= String "array", "items" .= sa ]
toBin [] = singleton 0
toBin as = toBin (length as) <> foldMap toBin as <> singleton 0
toAvron = Array . V.fromList . map toAvron
-- This is not suitable for incremental processing.
fromBin = get_blocks []
where
get_blocks acc = zagInt >>= \l -> if l == 0 then return $ reverse acc
else get_block acc l >>= get_blocks
get_block acc l = if l == 0 then return acc
else fromBin >>= \a -> get_block (a:acc) (l-1)
-- | A generic vector becomes an Avro array
instance Avro a => Avro (V.Vector a) where
toSchema as = do sa <- toSchema (V.head as)
return $ object [ "type" .= String "array", "items" .= sa ]
toBin as | V.null as = singleton 0
| otherwise = toBin (V.length as) <> foldMap toBin as <> singleton 0
toAvron = Array . V.map toAvron
-- This is not suitable for incremental processing.
fromBin = get_blocks []
where
get_blocks acc = zagInt >>= \l -> if l == 0 then return $ V.concat $ reverse acc
else get_block [] l >>=
get_blocks . (: acc) . V.fromListN l . reverse
get_block acc l = if l == 0 then return acc
else fromBin >>= \a -> get_block (a:acc) (l-1)
-- | An unboxed vector becomes an Avro array
instance (Avro a, U.Unbox a) => Avro (U.Vector a) where
toSchema as = do sa <- toSchema (U.head as)
return $ object [ "type" .= String "array", "items" .= sa ]
toBin as | U.null as = singleton 0
| otherwise = toBin (U.length as) <> U.foldr ((<>) . toBin) mempty as <> singleton 0
toAvron = Array . V.map toAvron . U.convert
-- This is not suitable for incremental processing.
fromBin = get_blocks []
where
get_blocks acc = zagInt >>= \l -> if l == 0 then return $ U.concat $ reverse acc
else get_block [] l >>=
get_blocks . (: acc) . U.fromListN l . reverse
get_block acc l = if l == 0 then return acc
else fromBin >>= \a -> get_block (a:acc) (l-1)
-- | A map from Text becomes an Avro map.
instance Avro a => Avro (H.HashMap T.Text a) where
toSchema m = do sa <- toSchema (m H.! T.empty)
return $ object [ "type" .= String "map", "values" .= sa ]
toBin as | H.null as = singleton 0
| otherwise = toBin (H.size as) <> H.foldrWithKey (\k v b -> toBin k <> toBin v <> b) (singleton 0) as
toAvron = Object . H.map toAvron
-- This is not suitable for incremental processing.
fromBin = get_blocks H.empty
where
get_blocks !acc = zagInt >>= \l -> if l == 0 then return acc
else get_block acc l >>= get_blocks
get_block !acc l = if l == 0 then return acc
else fromBin >>= \k -> fromBin >>= \v -> get_block (H.insert k v acc) (l-1)
-- * Some(!) complex types.
--
-- Enums: Enumerated symbols. This is generated automatically for sums
-- of empty alternatives. Constructor names become enum symbols.
-- Records: This is generated automatically for product types using
-- Haskell record syntax.
--
-- Unions: For Haskell sum-of-product types using record syntax for
-- every arm, an Avro instance resolving to a union of record can be
-- generated automatically. The constructor names become record type
-- names, their fields become record fields.
-- XXX Sometimes we build sum types containing sum types, Maybe being the
-- most obvious example. A (Maybe a) where a itself yields a union,
-- should probably yield a union with one more alternative (the null).
deriveAvros :: [Name] -> Q [Dec]
deriveAvros = liftM concat . mapM deriveAvro
deriveAvro :: Name -> Q [Dec]
deriveAvro nm = reify nm >>= case_info
where
err m = fail $ "cannot derive Avro for " ++ show nm ++ ", " ++ m
case_info (TyConI dec) = case_dec dec
case_info _ = err "it is not a type constructor"
simple_cons (NormalC _ []) = True
simple_cons _ = False
record_cons (RecC _ _) = True
record_cons _ = False
case_dec (NewtypeD _cxt _name _tyvarbndrs _con _) = err $ "don't know what to do for NewtypeD"
case_dec (DataD _cxt _name _tyvarbndrs cons _)
| all simple_cons cons = mk_enum_inst [ nm1 | NormalC nm1 [] <- cons ]
| all record_cons cons = mk_record_inst [ (nm1, vsts) | RecC nm1 vsts <- cons ]
| otherwise = err $ "don't know how to make an instance with these constructors"
case_dec _ = fail $ "is not a data or newtype declaration"
tolit = litE . StringL . nameBase
tolitlist (x:xs) = [| T.pack $(tolit x) : $(tolitlist xs) |]
tolitlist [ ] = [| [] |]
-- enum instance from list of names
mk_enum_inst :: [Name] -> Q [Dec]
mk_enum_inst nms =
[d| instance Avro $(conT nm) where
toSchema _ = return $ object [ "type" .= string "enum"
, "name" .= string $(tolit nm)
, "symbols" .= $(tolitlist nms) ]
toBin x = $(
return $ CaseE (VarE 'x)
[ Match (ConP nm1 [])
(NormalB (AppE (VarE 'zigInt)
(LitE (IntegerL i)))) []
| (i,nm1) <- zip [0..] nms ] )
fromBin = zagInt >>= \x -> $(
return $ CaseE (VarE 'x)
[ Match (LitP (IntegerL i))
(NormalB (AppE (VarE 'return)
(ConE nm1))) []
| (i,nm1) <- zip [0..] nms ] )
toAvron x = $(
return $ CaseE (VarE 'x)
[ Match (ConP nm1 [])
(NormalB (AppE (VarE 'string)
(LitE (StringL (nameBase nm1))))) []
| nm1 <- nms ] )
|]
-- record instance from record-like constructors
-- XXX maybe allow empty "normal" constructors, too
mk_record_inst :: [ (Name, [(Name, Strict, Type)]) ] -> Q [Dec]
mk_record_inst [(nm1,fs1)] =
[d| instance Avro $(conT nm) where
toSchema _ = $(mk_product_schema nm1 fs1)
toBin = $(to_bin_product fs1)
fromBin = $(from_bin_product [| return $(conE nm1) |] fs1)
toAvron = $(to_avron_product fs1)
|]
mk_record_inst arms =
[d| instance Avro $(conT nm) where
toSchema _ = Array . V.fromList <$> sequence
$( foldr (\(nm1,fs) k -> [| $(mk_product_schema nm1 fs) : $k |])
[| [] |] arms )
toBin =
$( do x <- newName "x"
LamE [VarP x] . CaseE (VarE x)
<$> sequence [ ($ []) . Match (RecP nm1 []) . NormalB
<$> [| zigInt $(litE (IntegerL i)) <> $(to_bin_product fs) $(varE x) |]
| (i,(nm1,fs)) <- zip [0..] arms ] )
fromBin = zagInt >>=
$( do x <- newName "x"
LamE [VarP x] . CaseE (VarE x)
<$> sequence [ ($ []) . Match (LitP (IntegerL i)) . NormalB
<$> from_bin_product [| return $(conE nm1) |] fs
| (i,(nm1,fs)) <- zip [0..] arms ] )
toAvron =
$( do x <- newName "x"
LamE [VarP x] . CaseE (VarE x)
<$> sequence [ ($ []) . Match (RecP nm1 []) . NormalB
<$> [| object [ $(tolit nm1) .= $(to_avron_product fs) $(varE x) ] |]
| (nm1,fs) <- arms ] )
|]
-- create schema for a product from a name and a list of fields
mk_product_schema nm1 tps =
[| $( fieldlist tps ) >>= \flds ->
memoObject $( tolit nm1 )
[ "type" .= string "record"
, "fields" .= Array (V.fromList flds) ] |]
fieldlist = foldr go [| return [] |]
where
go (nm1,_,tp) k =
[| do sch <- toSchema $(sigE (varE 'undefined) (return tp))
obs <- $k
return $ object [ "name" .= string $(tolit nm1)
, "type" .= sch ]
: obs |]
-- binary encoding of records: field by field.
to_bin_product nms =
[| \x -> $( foldr (\(nm1,_,_) k -> [| mappend (toBin ($(varE nm1) x)) $k |] )
[| mempty |] nms ) |]
from_bin_product =
foldl (\expr (_,_,_) -> [| $expr <*> fromBin |])
-- json encoding of records: fields in an object
to_avron_product nms =
[| \x -> object $(
foldr (\(nm1,_,_) k -> [| ($(tolit nm1) .= toAvron ($(varE nm1) x)) : $k |] )
[| [] |] nms ) |]
data ContainerOpts = ContainerOpts { objects_per_block :: Int
, filetype_label :: B.ByteString }
-- Writing a container file. This is an 'Enumeratee', we read a list of
-- suitable types, we write a header containing the generated schema,
-- and a series of blocks with serialized data.
writeAvroContainer :: (MonadIO m, Nullable s, ListLike s a, Avro a)
=> ContainerOpts -> Enumeratee s B.ByteString m r
writeAvroContainer ContainerOpts{..} out = do
ma <- peekStream
sync_marker <- liftIO $ B.pack <$> replicateM 16 randomIO
let schema = encode . runMkSchema . toSchema . fromJust $ ma
meta :: H.HashMap T.Text B.ByteString
meta = H.fromList [( "avro.schema", B.concat $ BL.toChunks schema )
,( "avro.codec", "null" )
,( "biohazard.filetype", filetype_label )]
hdr = fromByteString "Obj\1" <> toBin meta <> fromByteString sync_marker
let enc_blocks = iterLoop $ \out' -> do (num,code) <- joinI $ takeStream objects_per_block $
foldStream (\(!n,c) o -> (n+1, c <> toBin o)) (0::Int,mempty)
let code1 = toLazyByteString code
block = toBin num <> toBin (BL.length code1) <>
fromLazyByteString code1 <> fromByteString sync_marker
lift (enumList (BL.toChunks $ toLazyByteString block) out')
lift (enumList (BL.toChunks $ toLazyByteString hdr) out) >>= enc_blocks
-- XXX Possible codecs: null, zlib, snappy, lzma; all missing
-- XXX Should check schema on reading.
readAvroContainer :: (Monad m, ListLike s a, Avro a) => Enumeratee B.ByteString s m r
readAvroContainer out = do
4 <- heads "Obj\1" -- enough magic?
meta <- iterGet (fromBin :: Get (H.HashMap T.Text B.ByteString))
sync_marker <- iGetString 16
flip iterLoop out $ \o -> do num <- iterGet zagInt
sz <- iterGet fromBin
o' <- joinI $ takeStream sz $ -- codec goes here
convStream (LL.singleton `liftM` iterGet fromBin) o
16 <- heads sync_marker
return o'
-- | Repeatedly apply an 'Iteratee' to a value until end of stream.
-- Returns the final value.
iterLoop :: (Nullable s, Monad m) => (a -> Iteratee s m a) -> a -> Iteratee s m a
iterLoop it a = do e <- isFinished
if e then return a
else it a >>= iterLoop it
iterGet :: Monad m => Get a -> Iteratee B.ByteString m a
iterGet = go . runGetIncremental
where
go (Fail _ _ err) = throwErr (iterStrExc err)
go (Done rest _ a) = idone a (Chunk rest)
go (Partial dec) = liftI $ \ck -> case ck of
Chunk s -> go (dec $ Just s)
EOF mx -> case dec Nothing of
Fail _ _ err -> throwErr (iterStrExc err)
Partial _ -> throwErr (iterStrExc "<partial>")
Done rest _ a | B.null rest -> idone a (EOF mx)
| otherwise -> idone a (Chunk rest)