content: {-# LANGUAGE OverloadedStrings #-}
  
  module Parquet where
  
  import qualified DataFrame as D
  
  import Data.Int
  import Data.Text (Text)
  import Data.Time
  import GHC.IO (unsafePerformIO)
  import Test.HUnit
  
  allTypes :: D.DataFrame
  allTypes =
      D.fromNamedColumns
          [ ("id", D.fromList [4 :: Int32, 5, 6, 7, 2, 3, 0, 1])
          , ("bool_col", D.fromList [True, False, True, False, True, False, True, False])
          , ("tinyint_col", D.fromList [0 :: Int32, 1, 0, 1, 0, 1, 0, 1])
          , ("smallint_col", D.fromList [0 :: Int32, 1, 0, 1, 0, 1, 0, 1])
          , ("int_col", D.fromList [0 :: Int32, 1, 0, 1, 0, 1, 0, 1])
          , ("bigint_col", D.fromList [0 :: Int64, 10, 0, 10, 0, 10, 0, 10])
          , ("float_col", D.fromList [0 :: Float, 1.1, 0, 1.1, 0, 1.1, 0, 1.1])
          , ("double_col", D.fromList [0 :: Double, 10.1, 0, 10.1, 0, 10.1, 0, 10.1])
          ,
              ( "date_string_col"
              , D.fromList
                  [ "03/01/09" :: Text
                  , "03/01/09"
                  , "04/01/09"
                  , "04/01/09"
                  , "02/01/09"
                  , "02/01/09"
                  , "01/01/09"
                  , "01/01/09"
                  ]
              )
          , ("string_col", D.fromList (take 8 (cycle ["0" :: Text, "1"])))
          ,
              ( "timestamp_col"
              , D.fromList
                  [ UTCTime (fromGregorian 2009 3 1) (secondsToDiffTime 0)
                  , UTCTime (fromGregorian 2009 3 1) (secondsToDiffTime 60)
                  , UTCTime (fromGregorian 2009 4 1) (secondsToDiffTime 0)
                  , UTCTime (fromGregorian 2009 4 1) (secondsToDiffTime 60)
                  , UTCTime (fromGregorian 2009 2 1) (secondsToDiffTime 0)
                  , UTCTime (fromGregorian 2009 2 1) (secondsToDiffTime 60)
                  , UTCTime (fromGregorian 2009 1 1) (secondsToDiffTime 0)
                  , UTCTime (fromGregorian 2009 1 1) (secondsToDiffTime 60)
                  ]
              )
          ]
  
  allTypesPlain :: Test
  allTypesPlain =
      TestCase
          ( assertEqual
              "allTypesPlain"
              allTypes
              (unsafePerformIO (D.readParquet "./tests/data/alltypes_plain.parquet"))
          )
  
  allTypesPlainSnappy :: Test
  allTypesPlainSnappy =
      TestCase
          ( assertEqual
              "allTypesPlainSnappy"
              (D.filter "id" (`elem` [6 :: Int32, 7]) allTypes)
              (unsafePerformIO (D.readParquet "./tests/data/alltypes_plain.snappy.parquet"))
          )
  
  allTypesDictionary :: Test
  allTypesDictionary =
      TestCase
          ( assertEqual
              "allTypesPlainSnappy"
              (D.filter "id" (`elem` [0 :: Int32, 1]) allTypes)
              (unsafePerformIO (D.readParquet "./tests/data/alltypes_dictionary.parquet"))
          )
  
  mtCarsDataset :: D.DataFrame
  mtCarsDataset =
      D.fromNamedColumns
          [
              ( "model"
              , D.fromList
                  [ "Mazda RX4" :: Text
                  , "Mazda RX4 Wag"
                  , "Datsun 710"
                  , "Hornet 4 Drive"
                  , "Hornet Sportabout"
                  , "Valiant"
                  , "Duster 360"
                  , "Merc 240D"
                  , "Merc 230"
                  , "Merc 280"
                  , "Merc 280C"
                  , "Merc 450SE"
                  , "Merc 450SL"
                  , "Merc 450SLC"
                  , "Cadillac Fleetwood"
                  , "Lincoln Continental"
                  , "Chrysler Imperial"
                  , "Fiat 128"
                  , "Honda Civic"
                  , "Toyota Corolla"
                  , "Toyota Corona"
                  , "Dodge Challenger"
                  , "AMC Javelin"
                  , "Camaro Z28"
                  , "Pontiac Firebird"
                  , "Fiat X1-9"
                  , "Porsche 914-2"
                  , "Lotus Europa"
                  , "Ford Pantera L"
                  , "Ferrari Dino"
                  , "Maserati Bora"
                  , "Volvo 142E"
                  ]
              )
          ,
              ( "mpg"
              , D.fromList
                  [ 21.0 :: Double
                  , 21.0
                  , 22.8
                  , 21.4
                  , 18.7
                  , 18.1
                  , 14.3
                  , 24.4
                  , 22.8
                  , 19.2
                  , 17.8
                  , 16.4
                  , 17.3
                  , 15.2
                  , 10.4
                  , 10.4
                  , 14.7
                  , 32.4
                  , 30.4
                  , 33.9
                  , 21.5
                  , 15.5
                  , 15.2
                  , 13.3
                  , 19.2
                  , 27.3
                  , 26.0
                  , 30.4
                  , 15.8
                  , 19.7
                  , 15.0
                  , 21.4
                  ]
              )
          ,
              ( "cyl"
              , D.fromList
                  [ 6 :: Int32
                  , 6
                  , 4
                  , 6
                  , 8
                  , 6
                  , 8
                  , 4
                  , 4
                  , 6
                  , 6
                  , 8
                  , 8
                  , 8
                  , 8
                  , 8
                  , 8
                  , 4
                  , 4
                  , 4
                  , 4
                  , 8
                  , 8
                  , 8
                  , 8
                  , 4
                  , 4
                  , 4
                  , 8
                  , 6
                  , 8
                  , 4
                  ]
              )
          ,
              ( "disp"
              , D.fromList
                  [ 160.0 :: Double
                  , 160.0
                  , 108.0
                  , 258.0
                  , 360.0
                  , 225.0
                  , 360.0
                  , 146.7
                  , 140.8
                  , 167.6
                  , 167.6
                  , 275.8
                  , 275.8
                  , 275.8
                  , 472.0
                  , 460.0
                  , 440.0
                  , 78.7
                  , 75.7
                  , 71.1
                  , 120.1
                  , 318.0
                  , 304.0
                  , 350.0
                  , 400.0
                  , 79.0
                  , 120.3
                  , 95.1
                  , 351.0
                  , 145.0
                  , 301.0
                  , 121.0
                  ]
              )
          ,
              ( "hp"
              , D.fromList
                  [ 110 :: Int32
                  , 110
                  , 93
                  , 110
                  , 175
                  , 105
                  , 245
                  , 62
                  , 95
                  , 123
                  , 123
                  , 180
                  , 180
                  , 180
                  , 205
                  , 215
                  , 230
                  , 66
                  , 52
                  , 65
                  , 97
                  , 150
                  , 150
                  , 245
                  , 175
                  , 66
                  , 91
                  , 113
                  , 264
                  , 175
                  , 335
                  , 109
                  ]
              )
          ,
              ( "drat"
              , D.fromList
                  [ 3.9 :: Double
                  , 3.9
                  , 3.85
                  , 3.08
                  , 3.15
                  , 2.76
                  , 3.21
                  , 3.69
                  , 3.92
                  , 3.92
                  , 3.92
                  , 3.07
                  , 3.07
                  , 3.07
                  , 2.93
                  , 3.0
                  , 3.23
                  , 4.08
                  , 4.93
                  , 4.22
                  , 3.7
                  , 2.76
                  , 3.15
                  , 3.73
                  , 3.08
                  , 4.08
                  , 4.43
                  , 3.77
                  , 4.22
                  , 3.62
                  , 3.54
                  , 4.11
                  ]
              )
          ,
              ( "wt"
              , D.fromList
                  [ 2.62 :: Double
                  , 2.875
                  , 2.32
                  , 3.215
                  , 3.44
                  , 3.46
                  , 3.57
                  , 3.19
                  , 3.15
                  , 3.44
                  , 3.44
                  , 4.07
                  , 3.73
                  , 3.78
                  , 5.25
                  , 5.424
                  , 5.345
                  , 2.2
                  , 1.615
                  , 1.835
                  , 2.465
                  , 3.52
                  , 3.435
                  , 3.84
                  , 3.845
                  , 1.935
                  , 2.14
                  , 1.513
                  , 3.17
                  , 2.77
                  , 3.57
                  , 2.78
                  ]
              )
          ,
              ( "qsec"
              , D.fromList
                  [ 16.46 :: Double
                  , 17.02
                  , 18.61
                  , 19.44
                  , 17.02
                  , 20.22
                  , 15.84
                  , 20.0
                  , 22.9
                  , 18.3
                  , 18.9
                  , 17.4
                  , 17.6
                  , 18.0
                  , 17.98
                  , 17.82
                  , 17.42
                  , 19.47
                  , 18.52
                  , 19.9
                  , 20.01
                  , 16.87
                  , 17.3
                  , 15.41
                  , 17.05
                  , 18.9
                  , 16.7
                  , 16.9
                  , 14.5
                  , 15.5
                  , 14.6
                  , 18.6
                  ]
              )
          ,
              ( "vs"
              , D.fromList
                  [ 0 :: Int32
                  , 0
                  , 1
                  , 1
                  , 0
                  , 1
                  , 0
                  , 1
                  , 1
                  , 1
                  , 1
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 1
                  , 1
                  , 1
                  , 1
                  , 0
                  , 0
                  , 0
                  , 0
                  , 1
                  , 0
                  , 1
                  , 0
                  , 0
                  , 0
                  , 1
                  ]
              )
          ,
              ( "am"
              , D.fromList
                  [ 1 :: Int32
                  , 1
                  , 1
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 1
                  , 1
                  , 1
                  , 0
                  , 0
                  , 0
                  , 0
                  , 0
                  , 1
                  , 1
                  , 1
                  , 1
                  , 1
                  , 1
                  , 1
                  ]
              )
          ,
              ( "gear"
              , D.fromList
                  [ 4 :: Int32
                  , 4
                  , 4
                  , 3
                  , 3
                  , 3
                  , 3
                  , 4
                  , 4
                  , 4
                  , 4
                  , 3
                  , 3
                  , 3
                  , 3
                  , 3
                  , 3
                  , 4
                  , 4
                  , 4
                  , 3
                  , 3
                  , 3
                  , 3
                  , 3
                  , 4
                  , 5
                  , 5
                  , 5
                  , 5
                  , 5
                  , 4
                  ]
              )
          ,
              ( "carb"
              , D.fromList
                  [ 4 :: Int32
                  , 4
                  , 1
                  , 1
                  , 2
                  , 1
                  , 4
                  , 2
                  , 2
                  , 4
                  , 4
                  , 3
                  , 3
                  , 3
                  , 4
                  , 4
                  , 4
                  , 1
                  , 2
                  , 1
                  , 1
                  , 2
                  , 2
                  , 4
                  , 2
                  , 1
                  , 2
                  , 2
                  , 4
                  , 6
                  , 8
                  , 2
                  ]
              )
          ]
  
  mtCars :: Test
  mtCars =
      TestCase
          ( assertEqual
              "mt_cars"
              mtCarsDataset
              (unsafePerformIO (D.readParquet "./data/mtcars.parquet"))
          )
  
  -- Uncomment to run parquet tests.
  -- Currently commented because they don't run with github CI
  tests :: [Test]
  tests = [] -- [allTypesPlain, allTypesPlainSnappy, allTypesDictionary, mtCars]
  