dataframe 0.2.0.1 → 0.2.0.2
raw patch · 27 files changed
+2495/−296 lines, 27 filesdep +filepathdep +snappydep +zstd
Dependencies added: filepath, snappy, zstd
Files
- README.md +5/−2
- benchmark/Main.hs +3/−3
- dataframe.cabal +18/−6
- src/DataFrame.hs +2/−0
- src/DataFrame/Errors.hs +27/−32
- src/DataFrame/IO/Parquet.hs +1554/−0
- src/DataFrame/Internal/Column.hs +199/−85
- src/DataFrame/Internal/DataFrame.hs +1/−1
- src/DataFrame/Internal/Expression.hs +5/−4
- src/DataFrame/Internal/Types.hs +1/−1
- src/DataFrame/Lazy.hs +3/−0
- src/DataFrame/Lazy/IO/CSV.hs +327/−0
- src/DataFrame/Lazy/Internal/DataFrame.hs +92/−0
- src/DataFrame/Operations/Aggregation.hs +19/−19
- src/DataFrame/Operations/Core.hs +34/−18
- src/DataFrame/Operations/Merge.hs +39/−0
- src/DataFrame/Operations/Statistics.hs +10/−8
- src/DataFrame/Operations/Subset.hs +9/−5
- src/DataFrame/Operations/Transformations.hs +47/−12
- tests/Main.hs +5/−5
- tests/Operations/Apply.hs +16/−16
- tests/Operations/Derive.hs +4/−4
- tests/Operations/Filter.hs +13/−13
- tests/Operations/GroupBy.hs +34/−34
- tests/Operations/InsertColumn.hs +17/−17
- tests/Operations/Sort.hs +7/−7
- tests/Operations/Take.hs +4/−4
README.md view
@@ -67,11 +67,14 @@ ### Visual example  +## Supported input formats+* CSV+* Apache Parquet (still buggy and experimental)+ ## Future work-* Apache arrow and Parquet compatability+* Apache arrow compatability * Integration with common data formats (currently only supports CSV) * Support windowed plotting (currently only supports ASCII plots)-* Create a lazy API that builds an execution graph instead of running eagerly (will be used to compute on files larger than RAM) ## Contributing * Please first submit an issue and we can discuss there.
benchmark/Main.hs view
@@ -13,9 +13,9 @@ ns <- VU.replicateM n (randomRIO (-20.0 :: Double, 20.0)) xs <- VU.replicateM n (randomRIO (-20.0 :: Double, 20.0)) ys <- VU.replicateM n (randomRIO (-20.0 :: Double, 20.0))- let df = D.fromList [("first", D.UnboxedColumn ns),- ("second", D.UnboxedColumn xs),- ("third", D.UnboxedColumn ys)]+ let df = D.fromNamedColumns [("first", D.UnboxedColumn ns),+ ("second", D.UnboxedColumn xs),+ ("third", D.UnboxedColumn ys)] print $ D.mean "first" df print $ D.variance "second" df
dataframe.cabal view
@@ -1,6 +1,6 @@ cabal-version: 2.4 name: dataframe-version: 0.2.0.1+version: 0.2.0.2 synopsis: An intuitive, dynamically-typed DataFrame library. @@ -22,7 +22,8 @@ location: https://github.com/mchav/dataframe library- exposed-modules: DataFrame+ exposed-modules: DataFrame,+ DataFrame.Lazy other-modules: DataFrame.Internal.Types, DataFrame.Internal.Expression, DataFrame.Internal.Parsing,@@ -33,6 +34,7 @@ DataFrame.Internal.Row, DataFrame.Errors, DataFrame.Operations.Core,+ DataFrame.Operations.Merge, DataFrame.Operations.Subset, DataFrame.Operations.Sorting, DataFrame.Operations.Statistics,@@ -40,19 +42,25 @@ DataFrame.Operations.Typing, DataFrame.Operations.Aggregation, DataFrame.Display.Terminal.Plot,- DataFrame.IO.CSV+ DataFrame.IO.CSV,+ DataFrame.IO.Parquet,+ DataFrame.Lazy.IO.CSV,+ DataFrame.Lazy.Internal.DataFrame build-depends: base >= 4.17.2.0 && < 4.22, array ^>= 0.5, attoparsec >= 0.12 && <= 0.14.4, bytestring >= 0.11 && <= 0.12.2.0, containers >= 0.6.7 && < 0.8, directory >= 1.3.0.0 && <= 1.3.9.0,+ filepath >= 1.0.0.0 && <= 1.5.4.0, hashable >= 1.2 && <= 1.5.0.0,+ snappy >= 0.2.0.0 && <= 0.2.0.4, statistics >= 0.16.2.1 && <= 0.16.3.0, text >= 2.0 && <= 2.1.2, time >= 1.12 && <= 1.14, vector ^>= 0.13,- vector-algorithms ^>= 0.9+ vector-algorithms ^>= 0.9,+ zstd >= 0.1.2.0 && <= 0.1.3.0 hs-source-dirs: src default-language: Haskell2010 @@ -76,7 +84,9 @@ DataFrame.Operations.Typing, DataFrame.Operations.Aggregation, DataFrame.Display.Terminal.Plot,- DataFrame.IO.CSV+ DataFrame.IO.CSV,+ DataFrame.IO.Parquet,+ DataFrame.Lazy.IO.CSV build-depends: base >= 4.17.2.0 && < 4.22, array ^>= 0.5, attoparsec >= 0.12 && <= 0.14.4,@@ -84,11 +94,13 @@ containers >= 0.6.7 && < 0.8, directory >= 1.3.0.0 && <= 1.3.9.0, hashable >= 1.2 && <= 1.5.0.0,+ snappy >= 0.2.0.0 && <= 0.2.0.4, statistics >= 0.16.2.1 && <= 0.16.3.0, text >= 2.0 && <= 2.1.2, time >= 1.12 && <= 1.14, vector ^>= 0.13,- vector-algorithms ^>= 0.9+ vector-algorithms ^>= 0.9,+ zstd >= 0.1.2.0 && <= 0.1.3.0 hs-source-dirs: app, src default-language: Haskell2010
src/DataFrame.hs view
@@ -20,6 +20,8 @@ import DataFrame.Operations.Aggregation as D import DataFrame.Display.Terminal.Plot as D import DataFrame.IO.CSV as D+-- Support for Parquet is still experimental+import DataFrame.IO.Parquet as D import Data.Function
src/DataFrame/Errors.hs view
@@ -10,30 +10,30 @@ import Control.Exception import Data.Array+import Data.Either import DataFrame.Display.Terminal.Colours import Data.Typeable (Typeable) import Type.Reflection (TypeRep) +data TypeErrorContext a b = MkTypeErrorContext+ { userType :: Either String (TypeRep a)+ , expectedType :: Either String (TypeRep b)+ , errorColumnName :: Maybe String+ , callingFunctionName :: Maybe String+ }+ data DataFrameException where TypeMismatchException :: forall a b. (Typeable a, Typeable b)- => TypeRep a -- ^ given type- -> TypeRep b -- ^ expected type- -> T.Text -- ^ column name- -> T.Text -- ^ call point+ => TypeErrorContext a b -> DataFrameException- TypeMismatchException' :: forall a . (Typeable a)- => TypeRep a -- ^ given type- -> String -- ^ expected type- -> T.Text -- ^ column name- -> T.Text -- ^ call point- -> DataFrameException ColumnNotFoundException :: T.Text -> T.Text -> [T.Text] -> DataFrameException deriving (Exception) instance Show DataFrameException where show :: DataFrameException -> String- show (TypeMismatchException a b columnName callPoint) = addCallPointInfo columnName (Just callPoint) (typeMismatchError a b)- show (TypeMismatchException' a b columnName callPoint) = addCallPointInfo columnName (Just callPoint) (typeMismatchError' (show a) b)+ show (TypeMismatchException context) = let+ errorString = typeMismatchError (either id show (userType context)) (either id show (expectedType context))+ in addCallPointInfo (errorColumnName context) (callingFunctionName context) errorString show (ColumnNotFoundException columnName callPoint availableColumns) = columnNotFound columnName callPoint availableColumns columnNotFound :: T.Text -> T.Text -> [T.Text] -> String@@ -47,51 +47,46 @@ ++ T.unpack (guessColumnName name columns) ++ "?\n\n" -typeMismatchError ::- Type.Reflection.TypeRep a ->- Type.Reflection.TypeRep b ->- String-typeMismatchError a b = typeMismatchError' (show a) (show b)--typeMismatchError' :: String -> String -> String-typeMismatchError' givenType expectedType =+typeMismatchError :: String -> String -> String+typeMismatchError givenType expectedType = red $ red "\n\n[Error]: Type Mismatch" ++ "\n\tWhile running your code I tried to " ++ "get a column of type: " ++ red (show givenType)- ++ " but column was of type: "+ ++ " but the column in the dataframe was actually of type: " ++ green (show expectedType) -addCallPointInfo :: T.Text -> Maybe T.Text -> String -> String-addCallPointInfo name (Just cp) err =+addCallPointInfo :: Maybe String -> Maybe String -> String -> String+addCallPointInfo (Just name) (Just cp) err = err ++ ( "\n\tThis happened when calling function "- ++ brightGreen (T.unpack cp)+ ++ brightGreen cp ++ " on the column "- ++ brightGreen (T.unpack name)+ ++ brightGreen name ++ "\n\n"- ++ typeAnnotationSuggestion (T.unpack cp)+ ++ typeAnnotationSuggestion cp )-addCallPointInfo name Nothing err =+addCallPointInfo Nothing (Just cp) err = err ++ "\n" ++ typeAnnotationSuggestion cp+addCallPointInfo (Just name) Nothing err = err ++ ( "\n\tOn the column "- ++ T.unpack name+ ++ name ++ "\n\n"- ++ typeAnnotationSuggestion "<function>" )+addCallPointInfo Nothing Nothing err = err typeAnnotationSuggestion :: String -> String typeAnnotationSuggestion cp = "\n\n\tTry adding a type at the end of the function e.g " ++ "change\n\t\t"- ++ red (cp ++ " arg1 arg2")+ ++ red (cp ++ " ...") ++ " to \n\t\t"- ++ green ("(" ++ cp ++ " arg1 arg2 :: <Type>)")+ ++ green ("(" ++ cp ++ " ... :: <Type>)") ++ "\n\tor add " ++ "{-# LANGUAGE TypeApplications #-} to the top of your " ++ "file then change the call to \n\t\t"- ++ brightGreen (cp ++ " @<Type> arg1 arg2")+ ++ brightGreen (cp ++ " @<Type> ....") guessColumnName :: T.Text -> [T.Text] -> T.Text guessColumnName userInput columns = case map (\k -> (editDistance userInput k, k)) columns of
+ src/DataFrame/IO/Parquet.hs view
@@ -0,0 +1,1554 @@+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE RecordWildCards #-}++module DataFrame.IO.Parquet (+ readParquet+ ) where++import qualified Codec.Compression.Snappy as Snappy+import Codec.Compression.Zstd.Streaming+import Control.Monad+import qualified Data.ByteString as BSO+import qualified Data.ByteString.Char8 as BS+import Data.Char+import Data.Foldable+import qualified Data.Vector.Unboxed as VU+import Data.IORef+import qualified Data.Map as M+import Data.Maybe+import qualified Data.Text as T+import DataFrame.Internal.DataFrame (DataFrame)+import qualified DataFrame.Internal.DataFrame as DI+import qualified DataFrame.Internal.Column as DI+import qualified DataFrame.Operations.Core as DI+import Foreign+import GHC.Float+import GHC.IO (unsafePerformIO)+import System.IO++footerSize :: Integer+footerSize = 8++data ParquetType+ = PBOOLEAN+ | PINT32+ | PINT64+ | PINT96+ | PFLOAT+ | PDOUBLE+ | PBYTE_ARRAY+ | PFIXED_LEN_BYTE_ARRAY+ | PARQUET_TYPE_UNKNOWN+ deriving (Show, Eq)++parquetTypeFromInt :: Int32 -> ParquetType+parquetTypeFromInt 0 = PBOOLEAN+parquetTypeFromInt 1 = PINT32+parquetTypeFromInt 2 = PINT64+parquetTypeFromInt 3 = PINT96+parquetTypeFromInt 4 = PFLOAT+parquetTypeFromInt 5 = PDOUBLE+parquetTypeFromInt 6 = PBYTE_ARRAY+parquetTypeFromInt 7 = PFIXED_LEN_BYTE_ARRAY+parquetTypeFromInt _ = PARQUET_TYPE_UNKNOWN++data ParquetEncoding+ = EPLAIN+ | EPLAIN_DICTIONARY+ | ERLE+ | EBIT_PACKED+ | EDELTA_BINARY_PACKED+ | EDELTA_LENGTH_BYTE_ARRAY+ | EDELTA_BYTE_ARRAY+ | ERLE_DICTIONARY+ | EBYTE_STREAM_SPLIT+ | PARQUET_ENCODING_UNKNOWN+ deriving (Show, Eq)++parquetEncodingFromInt :: Int32 -> ParquetEncoding+parquetEncodingFromInt 0 = EPLAIN+parquetEncodingFromInt 2 = EPLAIN_DICTIONARY+parquetEncodingFromInt 3 = ERLE+parquetEncodingFromInt 4 = EBIT_PACKED+parquetEncodingFromInt 5 = EDELTA_BINARY_PACKED+parquetEncodingFromInt 6 = EDELTA_LENGTH_BYTE_ARRAY+parquetEncodingFromInt 7 = EDELTA_BYTE_ARRAY+parquetEncodingFromInt 8 = ERLE_DICTIONARY+parquetEncodingFromInt 9 = EBYTE_STREAM_SPLIT+parquetEncodingFromInt _ = PARQUET_ENCODING_UNKNOWN++data CompressionCodec+ = UNCOMPRESSED+ | SNAPPY+ | GZIP+ | LZO+ | BROTLI+ | LZ4+ | ZSTD+ | LZ4_RAW+ | COMPRESSION_CODEC_UNKNOWN+ deriving (Show, Eq)++compressionCodecFromInt :: Int32 -> CompressionCodec+compressionCodecFromInt 0 = UNCOMPRESSED+compressionCodecFromInt 1 = SNAPPY+compressionCodecFromInt 2 = GZIP+compressionCodecFromInt 3 = LZO+compressionCodecFromInt 4 = BROTLI+compressionCodecFromInt 5 = LZ4+compressionCodecFromInt 6 = ZSTD+compressionCodecFromInt 7 = LZ4_RAW+compressionCodecFromInt _ = COMPRESSION_CODEC_UNKNOWN++data ColumnStatistics = ColumnStatistics+ { columnMin :: [Word8],+ columnMax :: [Word8],+ columnNullCount :: Int64,+ columnDistictCount :: Int64,+ columnMinValue :: [Word8],+ columnMaxValue :: [Word8],+ isColumnMaxValueExact :: Bool,+ isColumnMinValueExact :: Bool+ }+ deriving (Show, Eq)++emptyColumnStatistics :: ColumnStatistics+emptyColumnStatistics = ColumnStatistics [] [] 0 0 [] [] False False++data PageType+ = DATA_PAGE+ | INDEX_PAGE+ | DICTIONARY_PAGE+ | DATA_PAGE_V2+ | PAGE_TYPE_UNKNOWN+ deriving (Show, Eq)++pageTypeFromInt :: Int32 -> PageType+pageTypeFromInt 0 = DATA_PAGE+pageTypeFromInt 1 = INDEX_PAGE+pageTypeFromInt 2 = DICTIONARY_PAGE+pageTypeFromInt 3 = DATA_PAGE_V2+pageTypeFromInt _ = PAGE_TYPE_UNKNOWN++data PageEncodingStats = PageEncodingStats+ { pageEncodingPageType :: PageType,+ pageEncoding :: ParquetEncoding,+ pagesWithEncoding :: Int32+ }+ deriving (Show, Eq)++emptyPageEncodingStats :: PageEncodingStats+emptyPageEncodingStats = PageEncodingStats PAGE_TYPE_UNKNOWN PARQUET_ENCODING_UNKNOWN 0++data SizeStatistics = SizeStatisics+ { unencodedByteArrayDataTypes :: Int64,+ repetitionLevelHistogram :: [Int64],+ definitionLevelHistogram :: [Int64]+ }+ deriving (Show, Eq)++emptySizeStatistics :: SizeStatistics+emptySizeStatistics = SizeStatisics 0 [] []++data BoundingBox = BoundingBox+ { xmin :: Double,+ xmax :: Double,+ ymin :: Double,+ ymax :: Double,+ zmin :: Double,+ zmax :: Double,+ mmin :: Double,+ mmax :: Double+ }+ deriving (Show, Eq)++emptyBoundingBox :: BoundingBox+emptyBoundingBox = BoundingBox 0 0 0 0 0 0 0 0++data GeospatialStatistics = GeospatialStatistics+ { bbox :: BoundingBox,+ geospatialTypes :: [Int32]+ }+ deriving (Show, Eq)++emptyGeospatialStatistics :: GeospatialStatistics+emptyGeospatialStatistics = GeospatialStatistics emptyBoundingBox []++emptyKeyValue :: KeyValue+emptyKeyValue = KeyValue {key = "", value = ""}++data ColumnMetaData = ColumnMetaData+ { columnType :: ParquetType,+ columnEncodings :: [ParquetEncoding],+ columnPathInSchema :: [String],+ columnCodec :: CompressionCodec,+ columnNumValues :: Int64,+ columnTotalUncompressedSize :: Int64,+ columnTotalCompressedSize :: Int64,+ columnKeyValueMetadata :: [KeyValue],+ columnDataPageOffset :: Int64,+ columnIndexPageOffset :: Int64,+ columnDictionaryPageOffset :: Int64,+ columnStatistics :: ColumnStatistics,+ columnEncodingStats :: [PageEncodingStats],+ bloomFilterOffset :: Int64,+ bloomFilterLength :: Int32,+ columnSizeStatistics :: SizeStatistics,+ columnGeospatialStatistics :: GeospatialStatistics+ }+ deriving (Show, Eq)++emptyColumnMetadata :: ColumnMetaData+emptyColumnMetadata = ColumnMetaData PARQUET_TYPE_UNKNOWN [] [] COMPRESSION_CODEC_UNKNOWN 0 0 0 [] 0 0 0 emptyColumnStatistics [] 0 0 emptySizeStatistics emptyGeospatialStatistics++data ColumnCryptoMetadata+ = COLUMN_CRYPTO_METADATA_UNKNOWN+ | ENCRYPTION_WITH_FOOTER_KEY+ | EncryptionWithColumnKey+ { columnCryptPathInSchema :: [String],+ columnKeyMetadata :: [Word8]+ }+ deriving (Show, Eq)++data ColumnChunk = ColumnChunk+ { columnChunkFilePath :: String,+ columnChunkMetadataFileOffset :: Int64,+ columnMetaData :: ColumnMetaData,+ columnChunkOffsetIndexOffset :: Int64,+ columnChunkOffsetIndexLength :: Int32,+ columnChunkColumnIndexOffset :: Int64,+ columnChunkColumnIndexLength :: Int32,+ cryptoMetadata :: ColumnCryptoMetadata,+ encryptedColumnMetadata :: [Word8]+ }+ deriving (Show, Eq)++emptyColumnChunk :: ColumnChunk+emptyColumnChunk = ColumnChunk "" 0 emptyColumnMetadata 0 0 0 0 COLUMN_CRYPTO_METADATA_UNKNOWN []++data SortingColumn = SortingColumn+ { columnIndex :: Int32,+ columnOrderDescending :: Bool,+ nullFirst :: Bool+ }+ deriving (Show, Eq)++emptySortingColumn :: SortingColumn+emptySortingColumn = SortingColumn 0 False False++data RowGroup = RowGroup+ { rowGroupColumns :: [ColumnChunk],+ totalByteSize :: Int64,+ rowGroupNumRows :: Int64,+ rowGroupSortingColumns :: [SortingColumn],+ fileOffset :: Int64,+ totalCompressedSize :: Int64,+ ordinal :: Int16+ }+ deriving (Show, Eq)++emptyRowGroup :: RowGroup+emptyRowGroup = RowGroup [] 0 0 [] 0 0 0++data ColumnOrder+ = TYPE_ORDER+ | COLUMN_ORDER_UNKNOWN+ deriving (Show, Eq)++data EncryptionAlgorithm+ = ENCRYPTION_ALGORITHM_UNKNOWN+ | AesGcmV1+ { aadPrefix :: [Word8],+ aadFileUnique :: [Word8],+ supplyAadPrefix :: Bool+ }+ | AesGcmCtrV1+ { aadPrefix :: [Word8],+ aadFileUnique :: [Word8],+ supplyAadPrefix :: Bool+ }+ deriving (Show, Eq)++data KeyValue = KeyValue+ { key :: String,+ value :: String+ }+ deriving (Show, Eq)++data FileMetadata = FileMetaData+ { version :: Int32,+ schema :: [SchemaElement],+ numRows :: Integer,+ rowGroups :: [RowGroup],+ keyValueMetadata :: [KeyValue],+ createdBy :: Maybe String,+ columnOrders :: [ColumnOrder],+ encryptionAlgorithm :: EncryptionAlgorithm,+ footerSigningKeyMetadata :: [Word8]+ }+ deriving (Show, Eq)++defaultMetadata :: FileMetadata+defaultMetadata =+ FileMetaData+ { version = 0,+ schema = [],+ numRows = 0,+ rowGroups = [],+ keyValueMetadata = [],+ createdBy = Nothing,+ columnOrders = [],+ encryptionAlgorithm = ENCRYPTION_ALGORITHM_UNKNOWN,+ footerSigningKeyMetadata = []+ }++readParquet :: String -> IO DataFrame+readParquet path = withBinaryFile path ReadMode $ \handle -> do+ (size, magicString) <- readMetadataSizeFromFooter handle+ when (magicString /= "PAR1") $ error "Invalid Parquet file"++ colMap <- newIORef (M.empty :: (M.Map T.Text DI.Column))+ colNames <- newIORef ([] :: [T.Text])++ fileMetadata <- readMetadata handle size+ forM_ (rowGroups fileMetadata) $ \r -> do+ forM_ (rowGroupColumns r) $ \c -> do+ let metadata = columnMetaData c+ let colDataPageOffset = columnDataPageOffset metadata+ let colDictionaryPageOffset = columnDictionaryPageOffset metadata+ let colStart = if colDictionaryPageOffset > 0 && colDataPageOffset > colDictionaryPageOffset+ then colDictionaryPageOffset+ else colDataPageOffset+ let colLength = columnTotalCompressedSize metadata+ columnBytes <-readBytes handle colStart colLength+ (maybePage, res) <- readPage (columnCodec metadata) columnBytes+ case maybePage of+ Just p -> if isDictionaryPage p+ then do+ (maybePage', res') <- readPage (columnCodec metadata) res+ let p' = fromMaybe (error "Empty page") maybePage'+ let schemaElem = filter (\se -> (elementName se) == (T.pack $ head (columnPathInSchema metadata))) (schema fileMetadata)+ let rep = if null schemaElem then UNKNOWN_REPETITION_TYPE else ((repetitionType . head) schemaElem)+ when (rep == REPEATED || rep == UNKNOWN_REPETITION_TYPE) (error $ "REPETITION TYPE NOT SUPPORTED: " ++ show rep)+ + case ((definitionLevelEncoding . pageTypeHeader . pageHeader ) p') of+ ERLE -> do+ let rleColumn = case columnType metadata of+ PBYTE_ARRAY -> readByteArrayColumn (pageBytes p)+ PDOUBLE -> readDoubleColumn (pageBytes p)+ PINT32 -> readInt32Column (pageBytes p)+ t -> error $ "UNKNOWN TYPE: " ++ (show t)+ let nbytes = littleEndianInt32 (take 4 (pageBytes p'))+ let rleDecoder = MkRleDecoder (drop 4 (pageBytes p')) 1 0 0++ -- Create index decoder+ let lvlByteLen = (fromIntegral nbytes + 4)+ let rleBytes = drop lvlByteLen (pageBytes p')+ let bitWidth = head rleBytes+ let indexDecoder = MkRleDecoder (tail rleBytes) (fromIntegral bitWidth) 0 0++ let finalCol = DI.takeColumn ((fromIntegral . dataPageHeaderNumValues . pageTypeHeader . pageHeader) p') (decodeDictionary rleColumn rleDecoder indexDecoder)+ let colName = T.pack $ head (columnPathInSchema metadata)++ modifyIORef' colNames (++[colName]) + modifyIORef' colMap (\m -> M.insertWith (\l r -> fromMaybe (error "UNEXPECTED") (DI.concatColumns l r)) colName finalCol m)+ other -> error $ "UNSUPPORTED ENCODING: " ++ (show other)+ else (error "PLAIN DATA PAGES NOT SUPPORTED")+ Nothing -> pure ()++ c' <- readIORef colMap+ colNames' <- readIORef colNames+ let asscList = map (\name -> (name, c' M.! name)) colNames'+ pure $ DI.fromNamedColumns asscList++decodeDictionary :: DI.Column -> RleDecoder -> RleDecoder -> DI.Column+decodeDictionary col rleDecoder indexDecoder+ | repCount indexDecoder > 0 = error "UNIMPLEMENTED: Repetition not supported"+ | litCount indexDecoder > 0 = decodeDictionary (DI.atIndicesStable (VU.map fromIntegral (getIndices indexDecoder)) col) rleDecoder (indexDecoder { litCount = 0 })+ | otherwise = let+ (finished, indexDecoder') = advance indexDecoder+ in if finished then col else decodeDictionary col rleDecoder indexDecoder'++advance :: RleDecoder -> (Bool, RleDecoder)+advance indexDecoder + | (rleDecoderData indexDecoder) == [] = (True, indexDecoder)+ | otherwise = let+ (indicator, remaining) = readUVarInt (rleDecoderData indexDecoder)+ isLiteral = (indicator .&. 1) /= 0+ countValues = (fromIntegral (indicator `shiftR` 1) :: Int32)+ litCount = if isLiteral then (countValues * 8) else 0+ in if isLiteral then (False, indexDecoder { rleDecoderData = remaining, litCount = litCount }) else (True, indexDecoder) -- (error "NON-LITERAL TYPES NOT YET SUPPORTED")++getIndices :: RleDecoder -> VU.Vector Word32+getIndices indexDecoder+ | rleBitWidth indexDecoder == 5 = unpackWidth5 (rleDecoderData indexDecoder)+ | rleBitWidth indexDecoder == 1 = unpackWidth1 (rleDecoderData indexDecoder)+ | rleBitWidth indexDecoder == 2 = unpackWidth2 (rleDecoderData indexDecoder)+ | rleBitWidth indexDecoder == 3 = unpackWidth3 (rleDecoderData indexDecoder)+ | otherwise = error $ "Unsupported bit width: " ++ (show (rleBitWidth indexDecoder))++unpackWidth5 :: [Word8] -> VU.Vector Word32+unpackWidth5 [] = VU.empty+unpackWidth5 bytes = let+ n0 = littleEndianWord32 $ take 4 bytes+ n1 = littleEndianWord32 $ take 4 $ drop 4 bytes+ n2 = littleEndianWord32 $ take 4 $ drop 8 bytes+ n3 = littleEndianWord32 $ take 4 $ drop 12 bytes+ n4 = littleEndianWord32 $ take 4 $ drop 16 bytes+ out0 = (n0 .>>. 0) `mod` (1 .<<. 5)+ out1 = (n0 .>>. 5) `mod` (1 .<<. 5)+ out2 = (n0 .>>. 10) `mod` (1 .<<. 5)+ out3 = (n0 .>>. 15) `mod` (1 .<<. 5)+ out4 = (n0 .>>. 20) `mod` (1 .<<. 5)+ out5 = (n0 .>>. 25) `mod` (1 .<<. 5)+ out6 = (n0 .>>. 30) .|. ((n1 `mod` (1 .<<. 3)) .<<. (5 - 3))+ out7 = (n1 .>>. 3) `mod` (1 .<<. 5)+ out8 = (n1 .>>. 8) `mod` (1 .<<. 5)+ out9 = (n1 .>>. 13) `mod` (1 .<<. 5)+ out10 = (n1 .>>. 18) `mod` (1 .<<. 5)+ out11 = (n1 .>>. 23) `mod` (1 .<<. 5)+ out12 = (n1 .>>. 28) .|. (n2 `mod` (1 .<<. 1)) .<<. (5 - 1)+ out13 = (n2 .>>. 1) `mod` (1 .<<. 5)+ out14 = (n2 .>>. 6) `mod` (1 .<<. 5)+ out15 = (n2 .>>. 11) `mod` (1 .<<. 5)+ out16 = (n2 .>>. 16) `mod` (1 .<<. 5)+ out17 = (n2 .>>. 21) `mod` (1 .<<. 5)+ out18 = (n2 .>>. 26) `mod` (1 .<<. 5)+ out19 = (n2 .>>. 31) .|. (n3 `mod` (1 .<<. 4)) .<<. (5 - 4)+ out20 = (n3 .>>. 4) `mod` (1 .<<. 5)+ out21 = (n3 .>>. 9) `mod` (1 .<<. 5)+ out22 = (n3 .>>. 14) `mod` (1 .<<. 5)+ out23 = (n3 .>>. 19) `mod` (1 .<<. 5)+ out24 = (n3 .>>. 24) `mod` (1 .<<. 5)+ out25 = (n3 .>>. 29) .|. (n4 `mod` (1 .<<. 2)) .<<. (5 - 2)+ out26 = (n4 .>>. 2) `mod` (1 .<<. 5)+ out27 = (n4 .>>. 7) `mod` (1 .<<. 5)+ out28 = (n4 .>>. 12) `mod` (1 .<<. 5)+ out29 = (n4 .>>. 17) `mod` (1 .<<. 5)+ out30 = (n4 .>>. 22) `mod` (1 .<<. 5)+ out31 = (n4 .>>. 27)+ in (VU.fromList [out0,out1,out2,out3,out4,out5,out6,out7,out8,out9,out10,out11,out12,out13,out14,out15,out16,out17,out18,out19,out20,out21,out22,out23,out24,out25,out26,out27,out28,out29,out30,out31]) VU.++ (unpackWidth5 (drop 20 bytes))++unpackWidth2, unpackWidth1, unpackWidth3 :: [Word8] -> VU.Vector Word32+unpackWidth1 [] = VU.empty+unpackWidth1 bytes = let+ n = littleEndianWord32 $ take 4 bytes+ in VU.fromList (map (\i -> (n .>>. i) .&. 1) [0..31]) VU.++ (unpackWidth1 (drop 4 bytes))+unpackWidth2 [] = VU.empty+unpackWidth2 bytes = let+ n = littleEndianWord32 $ take 4 bytes+ in VU.fromList (map (\i -> (n .>>. (i * 2)) .&. 1) [0..14] ++ [n .>>. 30]) VU.++ (unpackWidth2 (drop 4 bytes))+unpackWidth3 [] = VU.empty+unpackWidth3 bytes = let+ n0 = littleEndianWord32 $ take 4 bytes+ n1 = littleEndianWord32 $ take 4 $ drop 4 bytes+ n2 = littleEndianWord32 $ take 4 $ drop 8 bytes+ out0 = (n0 .>>. 0) `mod` (1 .<<. 3)+ out1 = (n0 .>>. 3) `mod` (1 .<<. 3)+ out2 = (n0 .>>. 6) `mod` (1 .<<. 3)+ out3 = (n0 .>>. 9) `mod` (1 .<<. 3)+ out4 = (n0 .>>. 12) `mod` (1 .<<. 3)+ out5 = (n0 .>>. 15) `mod` (1 .<<. 3)+ out6 = (n0 .>>. 18) `mod` (1 .<<. 3)+ out7 = (n0 .>>. 21) `mod` (1 .<<. 3)+ out8 = (n0 .>>. 24) `mod` (1 .<<. 3)+ out9 = (n0 .>>. 27) `mod` (1 .<<. 3)+ out10 = (n0 .>>. 30) .|. (n1 `mod` (1 .<<. 1)) .<<. (3 - 1)+ out11 = (n1 .>>. 1) `mod` (1 .<<. 3)+ out12 = (n1 .>>. 4) `mod` (1 .<<. 3)+ out13 = (n1 .>>. 7) `mod` (1 .<<. 3)+ out14 = (n1 .>>. 10) `mod` (1 .<<. 3)+ out15 = (n1 .>>. 13) `mod` (1 .<<. 3)+ out16 = (n1 .>>. 16) `mod` (1 .<<. 3)+ out17 = (n1 .>>. 19) `mod` (1 .<<. 3)+ out18 = (n1 .>>. 22) `mod` (1 .<<. 3)+ out19 = (n1 .>>. 25) `mod` (1 .<<. 3)+ out20 = (n1 .>>. 28) `mod` (1 .<<. 3)+ out21 = ((n1 .>>. 31) `mod` (1 .<<. 3)) .|. (n2 `mod` (1 .<<. 2)) .<<. (3 - 2)+ out22 = (n2 .>>. 2) `mod` (1 .<<. 3)+ out23 = (n2 .>>. 5) `mod` (1 .<<. 3)+ out24 = (n2 .>>. 8) `mod` (1 .<<. 3)+ out25 = (n2 .>>. 11) `mod` (1 .<<. 3)+ out26 = (n2 .>>. 14) `mod` (1 .<<. 3)+ out27 = (n2 .>>. 17) `mod` (1 .<<. 3)+ out28 = (n2 .>>. 20) `mod` (1 .<<. 3)+ out29 = (n2 .>>. 23) `mod` (1 .<<. 3)+ out30 = (n2 .>>. 26) `mod` (1 .<<. 3)+ out31 = (n2 .>>. 29)+ in (VU.fromList [out0,out1,out2,out3,out4,out5,out6,out7,out8,out9,out10,out11,out12,out13,out14,out15,out16,out17,out18,out19,out20,out21,out22,out23,out24,out25,out26,out27,out28,out29,out30,out31]) VU.++ (unpackWidth3 (drop 12 bytes))++data RleDecoder = MkRleDecoder { rleDecoderData :: [Word8]+ , rleBitWidth :: Int32+ , repCount :: Int32+ , litCount :: Int32+ } deriving (Show, Eq)++expandDictionary :: [Word8] -> [Word8]+expandDictionary (bitWidth:rest) = rest++readInt32Column :: [Word8] -> DI.Column+readInt32Column = DI.fromList . readPageInt32++readDoubleColumn :: [Word8] -> DI.Column+readDoubleColumn = DI.fromList . readPageWord64++readByteArrayColumn :: [Word8] -> DI.Column+readByteArrayColumn = DI.fromList .readPageBytes++readPageInt32 :: [Word8] -> [Int32]+readPageInt32 [] = []+readPageInt32 xs = (fromIntegral (littleEndianInt32 (take 4 xs))) : readPageInt32 (drop 4 xs)++readPageWord64 :: [Word8] -> [Double]+readPageWord64 [] = []+readPageWord64 xs = (castWord64ToDouble (littleEndianWord64 (take 8 xs))) : readPageWord64 (drop 8 xs)++readPageBytes :: [Word8] -> [T.Text]+readPageBytes [] = []+readPageBytes xs = let+ lenBytes = fromIntegral (littleEndianWord8 $ take 4 xs)+ totalBytesRead = lenBytes + 4+ in T.pack (map (chr . fromIntegral) $ take lenBytes (drop 4 xs)) : readPageBytes (drop totalBytesRead xs)++readPage :: CompressionCodec -> [Word8] -> IO (Maybe Page, [Word8])+readPage c [] = pure (Nothing, [])+readPage c columnBytes = do+ let (hdr, rem) = readPageHeader emptyPageHeader columnBytes 0+ let compressed = take (fromIntegral $ compressedPageSize hdr) rem+ + -- Weird round about way to uncompress zstd files compressed using the+ -- streaming API+ fullData <- case c of+ ZSTD -> do+ Consume dFunc <- decompress+ Consume dFunc' <- dFunc (BSO.pack compressed)+ Done res <- dFunc' BSO.empty+ pure res+ SNAPPY -> pure $ Snappy.decompress (BSO.pack compressed)+ UNCOMPRESSED -> pure (BSO.pack compressed)+ comp -> error ("UNSUPPORTED_COMPRESSION TYPE: " ++ (show comp))+ pure $ (Just $ Page hdr (BSO.unpack fullData), drop (fromIntegral $ compressedPageSize hdr) rem)++data Page = Page { pageHeader :: PageHeader+ , pageBytes :: [Word8] } deriving (Show, Eq)++data PageHeader = PageHeader { pageHeaderPageType :: PageType+ , uncompressedPageSize :: Int32+ , compressedPageSize ::Int32+ , pageHeaderCrcChecksum :: Int32+ , pageTypeHeader :: PageTypeHeader+ } deriving (Show, Eq)+++emptyPageHeader = PageHeader PAGE_TYPE_UNKNOWN 0 0 0 PAGE_TYPE_HEADER_UNKNOWN++isDataPage :: Page -> Bool+isDataPage page = case pageTypeHeader (pageHeader page) of+ DataPageHeader {..} -> True+ DataPageHeaderV2 {..} -> True+ _ -> False++isDictionaryPage :: Page -> Bool+isDictionaryPage page = case pageTypeHeader (pageHeader page) of+ DictionaryPageHeader {..} -> True+ _ -> False++data PageTypeHeader = DataPageHeader { dataPageHeaderNumValues :: Int32+ , dataPageHeaderEncoding :: ParquetEncoding+ , definitionLevelEncoding :: ParquetEncoding+ , repetitionLevelEncoding :: ParquetEncoding+ , dataPageHeaderStatistics :: ColumnStatistics+ }+ | DataPageHeaderV2 { dataPageHeaderV2NumValues :: Int32+ , dataPageHeaderV2NumNulls :: Int32+ , dataPageHeaderV2NumRows :: Int32+ , dataPageHeaderV2Encoding :: ParquetEncoding+ , definitionLevelByteLength :: Int32+ , repetitionLevelByteLength :: Int32+ , dataPageHeaderV2IsCompressed :: Bool+ , dataPageHeaderV2Statistics :: ColumnStatistics+ }+ | DictionaryPageHeader { dictionaryPageHeaderNumValues :: Int32+ , dictionaryPageHeaderEncoding :: ParquetEncoding+ , dictionaryPageIsSorted :: Bool + }+ | INDEX_PAGE_HEADER+ | PAGE_TYPE_HEADER_UNKNOWN deriving (Show, Eq)++emptyDictionaryPageHeader = DictionaryPageHeader 0 PARQUET_ENCODING_UNKNOWN False+emptyDataPageHeader = DataPageHeader 0 PARQUET_ENCODING_UNKNOWN PARQUET_ENCODING_UNKNOWN PARQUET_ENCODING_UNKNOWN emptyColumnStatistics+emptyDataPageHeaderV2 = DataPageHeaderV2 0 0 0 PARQUET_ENCODING_UNKNOWN 0 0 False emptyColumnStatistics ++readPageHeader :: PageHeader -> [Word8] -> Int16 -> (PageHeader, [Word8])+readPageHeader hdr [] _ = (hdr, [])+readPageHeader hdr xs lastFieldId = let+ fieldContents = readField' xs lastFieldId+ in case fieldContents of+ Nothing -> (hdr, tail xs)+ Just (rem, elemType, identifier) -> case identifier of+ 1 -> let+ (pType, rem') = readInt32FromBytes rem+ in readPageHeader (hdr {pageHeaderPageType = pageTypeFromInt pType}) rem' identifier+ 2 -> let+ (uncompressedPageSize, rem') = readInt32FromBytes rem+ in readPageHeader (hdr {uncompressedPageSize = uncompressedPageSize}) rem' identifier+ 3 -> let+ (compressedPageSize, rem') = readInt32FromBytes rem+ in readPageHeader (hdr {compressedPageSize = compressedPageSize}) rem' identifier+ 5 -> let+ (dataPageHeader, rem') = readPageTypeHeader emptyDataPageHeader rem 0+ in readPageHeader (hdr {pageTypeHeader = dataPageHeader}) rem' identifier+ 7 -> let+ (dictionaryPageHeader, rem') = readPageTypeHeader emptyDictionaryPageHeader rem 0+ in readPageHeader (hdr {pageTypeHeader = dictionaryPageHeader}) rem' identifier+ n -> error $ show n++readPageTypeHeader :: PageTypeHeader -> [Word8] -> Int16 -> (PageTypeHeader, [Word8])+readPageTypeHeader hdr [] _ = (hdr, [])+readPageTypeHeader hdr@(DictionaryPageHeader {..}) xs lastFieldId = let+ fieldContents = readField' xs lastFieldId+ in case fieldContents of+ Nothing -> (hdr, tail xs)+ Just (rem, elemType, identifier) -> case identifier of+ 1 -> let+ (numValues, rem') = readInt32FromBytes rem+ in readPageTypeHeader (hdr {dictionaryPageHeaderNumValues = numValues}) rem' identifier+ 2 -> let+ (enc, rem') = readInt32FromBytes rem+ in readPageTypeHeader (hdr {dictionaryPageHeaderEncoding = parquetEncodingFromInt enc}) rem' identifier+ 3 -> let+ (isSorted: rem') = rem+ in readPageTypeHeader (hdr {dictionaryPageIsSorted = isSorted == compactBooleanTrue}) rem' identifier+ n -> error $ show n+readPageTypeHeader hdr@(DataPageHeader {..}) xs lastFieldId = let+ fieldContents = readField' xs lastFieldId+ in case fieldContents of+ Nothing -> (hdr, tail xs)+ Just (rem, elemType, identifier) -> case identifier of+ 1 -> let+ (numValues, rem') = readInt32FromBytes rem+ in readPageTypeHeader (hdr {dataPageHeaderNumValues = numValues}) rem' identifier+ 2 -> let+ (enc, rem') = readInt32FromBytes rem+ in readPageTypeHeader (hdr {dataPageHeaderEncoding = parquetEncodingFromInt enc}) rem' identifier+ 3 -> let+ (enc, rem') = readInt32FromBytes rem+ in readPageTypeHeader (hdr {definitionLevelEncoding = parquetEncodingFromInt enc}) rem' identifier+ 4 -> let+ (enc, rem') = readInt32FromBytes rem+ in readPageTypeHeader (hdr {repetitionLevelEncoding = parquetEncodingFromInt enc}) rem' identifier+ 5 -> let+ (stats, rem') = readStatisticsFromBytes emptyColumnStatistics rem 0+ in readPageTypeHeader (hdr {dataPageHeaderStatistics = stats}) rem' identifier+ n -> error $ show n++readStatisticsFromBytes :: ColumnStatistics -> [Word8] -> Int16 -> (ColumnStatistics, [Word8])+readStatisticsFromBytes cs xs lastFieldId = let+ fieldContents = readField' xs lastFieldId+ in case fieldContents of+ Nothing -> (cs, tail xs)+ Just (rem, elemType, identifier) -> case identifier of+ 1 -> let+ (maxInBytes, rem') = readByteStringFromBytes rem+ in readStatisticsFromBytes (cs {columnMax = maxInBytes}) rem' identifier+ 2 -> let+ (minInBytes, rem') = readByteStringFromBytes rem+ in readStatisticsFromBytes (cs {columnMin = minInBytes}) rem' identifier+ 3 -> let+ (nullCount, rem') = readIntFromBytes @Int64 rem+ in readStatisticsFromBytes (cs {columnNullCount = nullCount}) rem' identifier+ 4 -> let+ (distinctCount, rem') = readIntFromBytes @Int64 rem+ in readStatisticsFromBytes (cs {columnDistictCount = distinctCount}) rem' identifier+ 5 -> let+ (maxInBytes, rem') = readByteStringFromBytes rem+ in readStatisticsFromBytes (cs {columnMaxValue = maxInBytes}) rem' identifier+ 6 -> let+ (minInBytes, rem') = readByteStringFromBytes rem+ in readStatisticsFromBytes (cs {columnMinValue = minInBytes}) rem' identifier+ 7 -> let+ (isMaxValueExact: rem') = rem+ in readStatisticsFromBytes (cs {isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue}) rem' identifier+ 8 -> let+ (isMinValueExact: rem') = rem+ in readStatisticsFromBytes (cs {isColumnMinValueExact = isMinValueExact == compactBooleanTrue}) rem' identifier+ n -> error $ show n++readBytes :: Handle -> Int64 -> Int64 -> IO [Word8]+readBytes handle colStart colLen = do+ buf <- mallocBytes (fromIntegral colLen) :: IO (Ptr Word8)+ hSeek handle AbsoluteSeek (fromIntegral colStart)+ _ <- hGetBuf handle buf (fromIntegral colLen) + columnBytes <- readByteString' buf colLen+ free buf+ pure columnBytes++numBytesInFile :: Handle -> IO Integer+numBytesInFile handle = do+ hSeek handle SeekFromEnd 0+ hTell handle++readMetadataSizeFromFooter :: Handle -> IO (Integer, BS.ByteString)+readMetadataSizeFromFooter handle = do+ footerOffSet <- numBytesInFile handle++ buf <- mallocBytes (fromIntegral footerSize) :: IO (Ptr Word8)++ hSeek handle AbsoluteSeek (fromIntegral $! footerOffSet - footerSize)+ n <- hGetBuf handle buf (fromIntegral footerSize)++ -- The bytes that store the metadata size.+ sizeBytes <- mapM (\i -> fromIntegral <$> (peekElemOff buf i :: IO Word8) :: IO Int32) [0 .. 3]+ let size = fromIntegral $ foldl' (.|.) 0 $! zipWith shift sizeBytes [0, 8, 16, 24]++ magicStringBytes <- mapM (\i -> peekElemOff buf i :: IO Word8) [4 .. 7]+ let magicString = BSO.pack magicStringBytes+ free buf+ return (size, magicString)++readMetadata :: Handle -> Integer -> IO FileMetadata+readMetadata handle size = do+ metaDataBuf <- mallocBytes (fromIntegral size) :: IO (Ptr Word8)+ footerOffSet <- numBytesInFile handle++ hSeek handle AbsoluteSeek (fromIntegral $! footerOffSet - footerSize - size)++ metadataBytesRead <- hGetBuf handle metaDataBuf (fromIntegral size)+ let lastFieldId = 0+ let fieldStack = []+ bufferPos <- newIORef (0 :: Int)+ metadata <- readFileMetaData defaultMetadata metaDataBuf bufferPos lastFieldId fieldStack+ free metaDataBuf+ return metadata++readFileMetaData :: FileMetadata -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO FileMetadata+readFileMetaData metadata metaDataBuf bufferPos lastFieldId fieldStack = do+ fieldContents <- readField metaDataBuf bufferPos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return metadata+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ version <- readIntFromBuffer @Int32 metaDataBuf bufferPos+ readFileMetaData (metadata {version = version}) metaDataBuf bufferPos identifier fieldStack+ 2 -> do+ -- We can do some type checking/exception handling here.+ -- Check elemType == List+ sizeAndType <- readAndAdvance bufferPos metaDataBuf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ -- type of the contents of the list.+ let elemType = toTType sizeAndType+ schemaElements <- replicateM sizeOnly (readSchemaElement defaultSchemaElement metaDataBuf bufferPos 0 [])+ readFileMetaData (metadata {schema = schemaElements}) metaDataBuf bufferPos identifier fieldStack+ 3 -> do+ numRows <- readIntFromBuffer @Int64 metaDataBuf bufferPos+ readFileMetaData (metadata {numRows = fromIntegral numRows}) metaDataBuf bufferPos identifier fieldStack+ 4 -> do+ -- We can do some type checking/exception handling here.+ -- Check elemType == List+ sizeAndType <- readAndAdvance bufferPos metaDataBuf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ -- type of the contents of the list.+ let elemType = toTType sizeAndType+ rowGroups <- replicateM sizeOnly (readRowGroup emptyRowGroup metaDataBuf bufferPos 0 [])+ readFileMetaData (metadata {rowGroups = rowGroups}) metaDataBuf bufferPos identifier fieldStack+ 5 -> do+ -- We can do some type checking/exception handling here.+ -- Check elemType == List+ sizeAndType <- readAndAdvance bufferPos metaDataBuf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ -- type of the contents of the list.+ let elemType = toTType sizeAndType+ keyValueMetadata <- replicateM sizeOnly (readKeyValue emptyKeyValue metaDataBuf bufferPos 0 [])+ readFileMetaData (metadata {keyValueMetadata = keyValueMetadata}) metaDataBuf bufferPos identifier fieldStack+ 6 -> do+ createdBy <- readString metaDataBuf bufferPos+ readFileMetaData (metadata {createdBy = Just createdBy}) metaDataBuf bufferPos identifier fieldStack+ 7 -> do+ sizeAndType <- readAndAdvance bufferPos metaDataBuf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ let elemType = toTType sizeAndType+ columnOrders <- replicateM sizeOnly (readColumnOrder metaDataBuf bufferPos 0 [])+ readFileMetaData (metadata {columnOrders = columnOrders}) metaDataBuf bufferPos identifier fieldStack+ 8 -> do+ encryptionAlgorithm <- readEncryptionAlgorithm metaDataBuf bufferPos 0 []+ readFileMetaData (metadata {encryptionAlgorithm = encryptionAlgorithm}) metaDataBuf bufferPos identifier fieldStack+ 9 -> do+ footerSigningKeyMetadata <- readByteString metaDataBuf bufferPos+ readFileMetaData (metadata {footerSigningKeyMetadata = footerSigningKeyMetadata}) metaDataBuf bufferPos identifier fieldStack+ n -> return $ error $ "UNIMPLEMENTED " ++ show n++readEncryptionAlgorithm :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm+readEncryptionAlgorithm buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return ENCRYPTION_ALGORITHM_UNKNOWN+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ readAesGcmV1 (AesGcmV1 {aadPrefix = [], aadFileUnique = [], supplyAadPrefix = False}) buf pos 0 []+ 2 -> do+ readAesGcmCtrV1 (AesGcmCtrV1 {aadPrefix = [], aadFileUnique = [], supplyAadPrefix = False}) buf pos 0 []+ n -> return ENCRYPTION_ALGORITHM_UNKNOWN++readAesGcmV1 :: EncryptionAlgorithm -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm+readAesGcmV1 v@(AesGcmV1 aadPrefix aadFileUnique supplyAadPrefix) buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return v+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ aadPrefix <- readByteString buf pos+ readAesGcmV1 (v {aadPrefix = aadPrefix}) buf pos lastFieldId fieldStack+ 2 -> do+ aadFileUnique <- readByteString buf pos+ readAesGcmV1 (v {aadFileUnique = aadFileUnique}) buf pos lastFieldId fieldStack+ 3 -> do+ supplyAadPrefix <- readAndAdvance pos buf+ readAesGcmV1 (v {supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) buf pos lastFieldId fieldStack+ _ -> return ENCRYPTION_ALGORITHM_UNKNOWN++readAesGcmCtrV1 :: EncryptionAlgorithm -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO EncryptionAlgorithm+readAesGcmCtrV1 v@(AesGcmCtrV1 aadPrefix aadFileUnique supplyAadPrefix) buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return v+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ aadPrefix <- readByteString buf pos+ readAesGcmCtrV1 (v {aadPrefix = aadPrefix}) buf pos lastFieldId fieldStack+ 2 -> do+ aadFileUnique <- readByteString buf pos+ readAesGcmCtrV1 (v {aadFileUnique = aadFileUnique}) buf pos lastFieldId fieldStack+ 3 -> do+ supplyAadPrefix <- readAndAdvance pos buf+ readAesGcmCtrV1 (v {supplyAadPrefix = supplyAadPrefix == compactBooleanTrue}) buf pos lastFieldId fieldStack+ _ -> return ENCRYPTION_ALGORITHM_UNKNOWN++readColumnOrder :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnOrder+readColumnOrder buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return COLUMN_ORDER_UNKNOWN+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ _ <- replicateM_ 2 (readTypeOrder buf pos 0 [])+ return TYPE_ORDER+ _ -> return COLUMN_ORDER_UNKNOWN++readTypeOrder :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnOrder+readTypeOrder buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return TYPE_ORDER+ Just (elemType, identifier) -> if elemType == STOP+ then return TYPE_ORDER + else readTypeOrder buf pos identifier fieldStack++data SchemaElement = SchemaElement+ { elementName :: T.Text,+ elementType :: TType,+ typeLength :: Int32,+ numChildren :: Int32,+ fieldId :: Int32,+ repetitionType :: RepetitionType,+ convertedType :: Int32,+ scale :: Int32,+ precision :: Int32,+ logicalType :: LogicalType+ }+ deriving (Show, Eq)++data RepetitionType = REQUIRED | OPTIONAL | REPEATED | UNKNOWN_REPETITION_TYPE deriving (Eq, Show)++data LogicalType+ = STRING_TYPE+ | MAP_TYPE+ | LIST_TYPE+ | ENUM_TYPE+ | DECIMAL_TYPE+ | DATE_TYPE+ | DecimalType {decimalTypePrecision :: Int32, decimalTypeScale :: Int32}+ | TimeType {isAdjustedToUTC :: Bool, unit :: TimeUnit}+ | -- This should probably have a different, more constrained TimeUnit type.+ TimestampType {isAdjustedToUTC :: Bool, unit :: TimeUnit}+ | IntType {bitWidth :: Int8, intIsSigned :: Bool}+ | LOGICAL_TYPE_UNKNOWN+ | JSON_TYPE+ | BSON_TYPE+ | UUID_TYPE+ | FLOAT16_TYPE+ | VariantType {specificationVersion :: Int8}+ | GeometryType {crs :: T.Text}+ | GeographyType {crs :: T.Text, algorithm :: EdgeInterpolationAlgorithm}+ deriving (Eq, Show)++data TimeUnit+ = MILLISECONDS+ | MICROSECONDS+ | NANOSECONDS+ | TIME_UNIT_UNKNOWN+ deriving (Eq, Show)++data EdgeInterpolationAlgorithm+ = SPHERICAL+ | VINCENTY+ | THOMAS+ | ANDOYER+ | KARNEY+ deriving (Eq, Show)++repetitionTypeFromInt :: Int32 -> RepetitionType+repetitionTypeFromInt 0 = REQUIRED+repetitionTypeFromInt 1 = OPTIONAL+repetitionTypeFromInt 2 = REPEATED+repetitionTypeFromInt _ = UNKNOWN_REPETITION_TYPE++defaultSchemaElement :: SchemaElement+defaultSchemaElement = SchemaElement "" STOP 0 0 (-1) UNKNOWN_REPETITION_TYPE 0 0 0 LOGICAL_TYPE_UNKNOWN++toIntegralType :: Int32 -> TType+toIntegralType n+ | n == 0 = BOOL+ | n == 1 = I32+ | n == 2 = I64+ | n == 3 = I64+ | n == 4 = DOUBLE+ | n == 5 = DOUBLE+ | n == 6 = STRING+ | n == 7 = STRING+ | otherwise = error $ "Unknown integral type: " ++ show n++readSchemaElement :: SchemaElement -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO SchemaElement+readSchemaElement schemaElement buf pos lastFieldId fieldStack = do+ t <- readAndAdvance pos buf+ if t .&. 0x0f == 0+ then return schemaElement+ else do+ let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16+ identifier <-+ if modifier == 0+ then readIntFromBuffer @Int16 buf pos+ else return (lastFieldId + modifier)+ let elemType = toTType (t .&. 0x0f)+ case identifier of+ 1 -> do+ schemaElemType <- toIntegralType <$> readInt32FromBuffer buf pos+ readSchemaElement (schemaElement {elementType = schemaElemType}) buf pos identifier fieldStack+ 2 -> do+ typeLength <- readInt32FromBuffer buf pos+ readSchemaElement (schemaElement {typeLength = typeLength}) buf pos identifier fieldStack+ 3 -> do+ fieldRepetitionType <- readInt32FromBuffer buf pos+ readSchemaElement (schemaElement {repetitionType = repetitionTypeFromInt fieldRepetitionType}) buf pos identifier fieldStack+ 4 -> do+ nameSize <- readVarIntFromBuffer @Int buf pos+ contents <- replicateM nameSize (readAndAdvance pos buf)+ readSchemaElement (schemaElement {elementName = T.pack (map (chr . fromIntegral) contents)}) buf pos identifier fieldStack+ 5 -> do+ numChildren <- readInt32FromBuffer buf pos+ readSchemaElement (schemaElement {numChildren = numChildren}) buf pos identifier fieldStack+ 6 -> do+ convertedType <- readInt32FromBuffer buf pos+ readSchemaElement (schemaElement {convertedType = convertedType}) buf pos identifier fieldStack+ 7 -> do+ scale <- readInt32FromBuffer buf pos+ readSchemaElement (schemaElement {scale = scale}) buf pos identifier fieldStack+ 8 -> do+ precision <- readInt32FromBuffer buf pos+ readSchemaElement (schemaElement {precision = precision}) buf pos identifier fieldStack+ 9 -> do+ fieldId <- readInt32FromBuffer buf pos+ readSchemaElement (schemaElement {fieldId = fieldId}) buf pos identifier fieldStack+ 10 -> do+ logicalType <- readLogicalType buf pos 0 []+ readSchemaElement (schemaElement {logicalType = logicalType}) buf pos identifier fieldStack+ _ -> error $ show identifier -- return schemaElement++readLogicalType :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType+readLogicalType buf pos lastFieldId fieldStack = do+ t <- readAndAdvance pos buf+ if t .&. 0x0f == 0+ then return LOGICAL_TYPE_UNKNOWN+ else do+ let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16+ identifier <-+ if modifier == 0+ then readIntFromBuffer @Int16 buf pos+ else return (lastFieldId + modifier)+ let elemType = toTType (t .&. 0x0f)+ case identifier of+ 1 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return STRING_TYPE+ 2 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return MAP_TYPE+ 3 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return LIST_TYPE+ 4 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return ENUM_TYPE+ 5 -> do+ _ <- readField buf pos 0 []+ readDecimalType (DecimalType {decimalTypeScale = 0, decimalTypePrecision = 0}) buf pos 0 []+ 6 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return DATE_TYPE+ 7 -> do+ _ <- readField buf pos 0 []+ readTimeType (TimeType {isAdjustedToUTC = False, unit = MILLISECONDS}) buf pos 0 []+ 8 -> do+ _ <- readField buf pos 0 []+ readTimeType (TimestampType {isAdjustedToUTC = False, unit = MILLISECONDS}) buf pos 0 []+ -- Apparently reserved for interval types+ 9 -> return LOGICAL_TYPE_UNKNOWN+ 10 -> do+ _ <- readField buf pos 0 []+ readIntType (IntType {intIsSigned = False, bitWidth = 0}) buf pos 0 []+ 11 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return LOGICAL_TYPE_UNKNOWN+ 12 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return JSON_TYPE+ 13 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return BSON_TYPE+ 14 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return UUID_TYPE+ 15 -> do+ replicateM_ 2 (readField buf pos 0 [])+ return FLOAT16_TYPE+ 16 -> do+ _ <- readField buf pos 0 []+ return VariantType {specificationVersion = 1}+ 17 -> do+ _ <- readField buf pos 0 []+ return GeometryType {crs = ""}+ 18 -> do+ _ <- readField buf pos 0 []+ return GeographyType {crs = "", algorithm = SPHERICAL}+ _ -> return LOGICAL_TYPE_UNKNOWN++readIntType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType+readIntType v@(IntType bitWidth intIsSigned) buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return v+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ bitWidth <- readAndAdvance pos buf+ readIntType (v {bitWidth = fromIntegral bitWidth}) buf pos lastFieldId fieldStack+ 2 -> do+ -- TODO: Check for empty+ intIsSigned <- readAndAdvance pos buf+ readIntType (v {intIsSigned = intIsSigned == compactBooleanTrue}) buf pos lastFieldId fieldStack+ _ -> error $ "UNKNOWN field ID for IntType" ++ show identifier++readDecimalType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType+readDecimalType v@(DecimalType p s) buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return v+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ s' <- readInt32FromBuffer buf pos+ readDecimalType (v {decimalTypeScale = s'}) buf pos lastFieldId fieldStack+ 2 -> do+ p' <- readInt32FromBuffer buf pos+ readDecimalType (v {decimalTypePrecision = p'}) buf pos lastFieldId fieldStack+ _ -> error $ "UNKNOWN field ID for DecimalType" ++ show identifier++readTimeType :: LogicalType -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO LogicalType+readTimeType v@(TimeType _ _) buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return v+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ -- TODO: Check for empty+ isAdjustedToUTC <- readAndAdvance pos buf+ readTimeType (v {isAdjustedToUTC = isAdjustedToUTC == compactBooleanTrue}) buf pos lastFieldId fieldStack+ 2 -> do+ u <- readUnit buf pos 0 []+ readTimeType (v {unit = u}) buf pos lastFieldId fieldStack+ _ -> error $ "UNKNOWN field ID for TimeType" ++ show identifier+readTimeType v@(TimestampType _ _) buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return v+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ -- TODO: Check for empty+ isAdjustedToUTC <- readAndAdvance pos buf+ readTimeType (v {isAdjustedToUTC = isAdjustedToUTC == compactBooleanTrue}) buf pos lastFieldId fieldStack+ 2 -> do+ u <- readUnit buf pos 0 []+ readTimeType (v {unit = u}) buf pos lastFieldId fieldStack+ _ -> error $ "UNKNOWN field ID for TimestampType" ++ show identifier++readUnit :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO TimeUnit+readUnit buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return TIME_UNIT_UNKNOWN+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ _ <- readField buf pos 0 []+ return MILLISECONDS+ 2 -> do+ _ <- readField buf pos 0 []+ return MICROSECONDS+ 3 -> do+ _ <- readField buf pos 0 []+ return NANOSECONDS+ _ -> return TIME_UNIT_UNKNOWN++readRowGroup :: RowGroup -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO RowGroup+readRowGroup r buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return r+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ sizeAndType <- readAndAdvance pos buf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ -- type of the contents of the list.+ let elemType = toTType sizeAndType+ columnChunks <- replicateM sizeOnly (readColumnChunk emptyColumnChunk buf pos 0 [])+ readRowGroup (r {rowGroupColumns = columnChunks}) buf pos identifier fieldStack+ 2 -> do+ totalBytes <- readIntFromBuffer @Int64 buf pos+ readRowGroup (r {totalByteSize = totalBytes}) buf pos identifier fieldStack+ 3 -> do+ nRows <- readIntFromBuffer @Int64 buf pos+ readRowGroup (r {rowGroupNumRows = nRows}) buf pos identifier fieldStack+ 4 -> return r+ 5 -> do+ offset <- readIntFromBuffer @Int64 buf pos+ readRowGroup (r {fileOffset = offset}) buf pos identifier fieldStack+ 6 -> do+ compressedSize <- readIntFromBuffer @Int64 buf pos+ readRowGroup (r {totalCompressedSize = compressedSize}) buf pos identifier fieldStack+ 7 -> do+ ordinal <- readIntFromBuffer @Int16 buf pos+ readRowGroup (r {ordinal = ordinal}) buf pos identifier fieldStack+ _ -> error $ "Unknown row group field: " ++ show identifier++readColumnChunk :: ColumnChunk -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnChunk+readColumnChunk c buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return c+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ stringSize <- readVarIntFromBuffer @Int buf pos+ contents <- map (chr . fromIntegral) <$> replicateM stringSize (readAndAdvance pos buf)+ readColumnChunk (c {columnChunkFilePath = contents}) buf pos identifier fieldStack+ 2 -> do+ columnChunkMetadataFileOffset <- readIntFromBuffer @Int64 buf pos+ readColumnChunk (c {columnChunkMetadataFileOffset = columnChunkMetadataFileOffset}) buf pos identifier fieldStack+ 3 -> do+ columnMetadata <- readColumnMetadata emptyColumnMetadata buf pos 0 []+ readColumnChunk (c {columnMetaData = columnMetadata}) buf pos identifier fieldStack + 4 -> do+ columnOffsetIndexOffset <- readIntFromBuffer @Int64 buf pos+ readColumnChunk (c {columnChunkOffsetIndexOffset = columnOffsetIndexOffset}) buf pos identifier fieldStack+ 5 -> do+ columnOffsetIndexLength <- readInt32FromBuffer buf pos+ readColumnChunk (c {columnChunkOffsetIndexLength = columnOffsetIndexLength}) buf pos identifier fieldStack+ 6 -> do+ columnChunkColumnIndexOffset <- readIntFromBuffer @Int64 buf pos+ readColumnChunk (c {columnChunkColumnIndexOffset = columnChunkColumnIndexOffset}) buf pos identifier fieldStack+ 7 -> do+ columnChunkColumnIndexLength <- readInt32FromBuffer buf pos+ readColumnChunk (c {columnChunkColumnIndexLength = columnChunkColumnIndexLength}) buf pos identifier fieldStack+ _ -> return c++readColumnMetadata :: ColumnMetaData -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnMetaData+readColumnMetadata cm buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return cm+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ cType <- parquetTypeFromInt <$> readInt32FromBuffer buf pos+ readColumnMetadata (cm {columnType = cType}) buf pos identifier []+ 2 -> do+ sizeAndType <- readAndAdvance pos buf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ let elemType = toTType sizeAndType+ encodings <- replicateM sizeOnly (readParquetEncoding buf pos 0 [])+ readColumnMetadata (cm {columnEncodings = encodings}) buf pos identifier fieldStack+ 3 -> do+ sizeAndType <- readAndAdvance pos buf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ let elemType = toTType sizeAndType+ paths <- replicateM sizeOnly (readString buf pos)+ readColumnMetadata (cm {columnPathInSchema = paths}) buf pos identifier fieldStack+ 4 -> do+ cType <- compressionCodecFromInt <$> readInt32FromBuffer buf pos+ readColumnMetadata (cm {columnCodec = cType}) buf pos identifier []+ 5 -> do+ numValues <- readIntFromBuffer @Int64 buf pos+ readColumnMetadata (cm {columnNumValues = numValues}) buf pos identifier []+ 6 -> do+ columnTotalUncompressedSize <- readIntFromBuffer @Int64 buf pos+ readColumnMetadata (cm {columnTotalUncompressedSize = columnTotalUncompressedSize}) buf pos identifier []+ 7 -> do+ columnTotalCompressedSize <- readIntFromBuffer @Int64 buf pos+ readColumnMetadata (cm {columnTotalCompressedSize = columnTotalCompressedSize}) buf pos identifier []+ 8 -> do+ sizeAndType <- readAndAdvance pos buf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ let elemType = toTType sizeAndType+ columnKeyValueMetadata <- replicateM sizeOnly (readKeyValue emptyKeyValue buf pos 0 [])+ readColumnMetadata (cm {columnKeyValueMetadata = columnKeyValueMetadata}) buf pos identifier fieldStack+ 9 -> do+ columnDataPageOffset <- readIntFromBuffer @Int64 buf pos+ readColumnMetadata (cm {columnDataPageOffset = columnDataPageOffset}) buf pos identifier []+ 10 -> do+ columnIndexPageOffset <- readIntFromBuffer @Int64 buf pos+ readColumnMetadata (cm {columnIndexPageOffset = columnIndexPageOffset}) buf pos identifier []+ 11 -> do+ columnDictionaryPageOffset <- readIntFromBuffer @Int64 buf pos+ readColumnMetadata (cm {columnDictionaryPageOffset = columnDictionaryPageOffset}) buf pos identifier []+ 12 -> do+ stats <- readStatistics emptyColumnStatistics buf pos 0 []+ readColumnMetadata (cm {columnStatistics = stats}) buf pos identifier fieldStack+ 13 -> do+ sizeAndType <- readAndAdvance pos buf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ let elemType = toTType sizeAndType+ pageEncodingStats <- replicateM sizeOnly (readPageEncodingStats emptyPageEncodingStats buf pos 0 [])+ readColumnMetadata (cm {columnEncodingStats = pageEncodingStats}) buf pos identifier fieldStack+ 14 -> do+ bloomFilterOffset <- readIntFromBuffer @Int64 buf pos+ readColumnMetadata (cm {bloomFilterOffset = bloomFilterOffset}) buf pos identifier []+ 15 -> do+ bloomFilterLength <- readInt32FromBuffer buf pos+ readColumnMetadata (cm {bloomFilterLength = bloomFilterLength}) buf pos identifier []+ 16 -> do+ stats <- readSizeStatistics emptySizeStatistics buf pos 0 []+ readColumnMetadata (cm {columnSizeStatistics = stats}) buf pos identifier fieldStack+ 17 -> return $ error "UNIMPLEMENTED"+ _ -> return cm++readParquetEncoding :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ParquetEncoding+readParquetEncoding buf pos lastFieldId fieldStack = parquetEncodingFromInt <$> readInt32FromBuffer buf pos++readPageEncodingStats :: PageEncodingStats -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO PageEncodingStats+readPageEncodingStats pes buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return pes+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ pType <- pageTypeFromInt <$> readInt32FromBuffer buf pos+ readPageEncodingStats (pes {pageEncodingPageType = pType}) buf pos identifier []+ 2 -> do+ pEnc <- parquetEncodingFromInt <$> readInt32FromBuffer buf pos+ readPageEncodingStats (pes {pageEncoding = pEnc}) buf pos identifier []+ 3 -> do+ encodedCount <- readInt32FromBuffer buf pos+ readPageEncodingStats (pes {pagesWithEncoding = encodedCount}) buf pos identifier []+ _ -> pure pes++readStatistics :: ColumnStatistics -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO ColumnStatistics+readStatistics cs buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return cs+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ maxInBytes <- readByteString buf pos+ readStatistics (cs {columnMax = maxInBytes}) buf pos identifier fieldStack+ 2 -> do+ minInBytes <- readByteString buf pos+ readStatistics (cs {columnMin = minInBytes}) buf pos identifier fieldStack+ 3 -> do+ nullCount <- readIntFromBuffer @Int64 buf pos+ readStatistics (cs {columnNullCount = nullCount}) buf pos identifier fieldStack+ 4 -> do+ distinctCount <- readIntFromBuffer @Int64 buf pos+ readStatistics (cs {columnDistictCount = distinctCount}) buf pos identifier fieldStack+ 5 -> do+ maxInBytes <- readByteString buf pos+ readStatistics (cs {columnMaxValue = maxInBytes}) buf pos identifier fieldStack+ 6 -> do+ minInBytes <- readByteString buf pos+ readStatistics (cs {columnMinValue = minInBytes}) buf pos identifier fieldStack+ 7 -> do+ isMaxValueExact <- readAndAdvance pos buf+ readStatistics (cs {isColumnMaxValueExact = isMaxValueExact == compactBooleanTrue}) buf pos identifier fieldStack+ 8 -> do+ isMinValueExact <- readAndAdvance pos buf+ readStatistics (cs {isColumnMinValueExact = isMinValueExact == compactBooleanTrue}) buf pos identifier fieldStack+ _ -> pure cs++readSizeStatistics :: SizeStatistics -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO SizeStatistics+readSizeStatistics ss buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return ss+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ unencodedByteArrayDataTypes <- readIntFromBuffer @Int64 buf pos+ readSizeStatistics (ss {unencodedByteArrayDataTypes = unencodedByteArrayDataTypes}) buf pos identifier fieldStack+ 2 -> do+ sizeAndType <- readAndAdvance pos buf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ let elemType = toTType sizeAndType+ repetitionLevelHistogram <- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos)+ readSizeStatistics (ss {repetitionLevelHistogram = repetitionLevelHistogram}) buf pos identifier fieldStack+ 3 -> do+ sizeAndType <- readAndAdvance pos buf+ let sizeOnly = fromIntegral ((sizeAndType `shiftR` 4) .&. 0x0f) :: Int+ let elemType = toTType sizeAndType+ definitionLevelHistogram <- replicateM sizeOnly (readIntFromBuffer @Int64 buf pos)+ readSizeStatistics (ss {definitionLevelHistogram = definitionLevelHistogram}) buf pos identifier fieldStack+ _ -> pure ss++readKeyValue :: KeyValue -> Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO KeyValue+readKeyValue kv buf pos lastFieldId fieldStack = do+ fieldContents <- readField buf pos lastFieldId fieldStack+ case fieldContents of+ Nothing -> return kv+ Just (elemType, identifier) -> case identifier of+ 1 -> do+ k <- readString buf pos+ readKeyValue (kv {key = k}) buf pos identifier fieldStack+ 2 -> do+ v <- readString buf pos+ readKeyValue (kv {key = v}) buf pos identifier fieldStack+ _ -> return kv++readString :: Ptr Word8 -> IORef Int -> IO String+readString buf pos = do+ nameSize <- readVarIntFromBuffer @Int buf pos+ map (chr . fromIntegral) <$> replicateM nameSize (readAndAdvance pos buf)++readByteStringFromBytes :: [Word8] -> ([Word8], [Word8])+readByteStringFromBytes xs = let+ (size, rem) = readVarIntFromBytes @Int xs+ in (take size rem, drop size rem)++readByteString :: Ptr Word8 -> IORef Int -> IO [Word8]+readByteString buf pos = do+ size <- readVarIntFromBuffer @Int buf pos+ replicateM size (readAndAdvance pos buf)++readByteString' :: Ptr Word8 -> Int64 -> IO [Word8]+readByteString' buf size = mapM (`readSingleByte` buf) [0..(size - 1)]++readField :: Ptr Word8 -> IORef Int -> Int16 -> [Int16] -> IO (Maybe (TType, Int16))+readField buf pos lastFieldId fieldStack = do+ t <- readAndAdvance pos buf+ if t .&. 0x0f == 0+ then return Nothing+ else do+ let modifier = fromIntegral ((t .&. 0xf0) `shiftR` 4) :: Int16+ identifier <-+ if modifier == 0+ then readIntFromBuffer @Int16 buf pos+ else return (lastFieldId + modifier)+ let elemType = toTType (t .&. 0x0f)+ pure $ Just (elemType, identifier)++readField' :: [Word8] -> Int16 -> Maybe ([Word8], TType, Int16)+readField' [] _ = Nothing+readField' (x:xs) lastFieldId+ | x .&. 0x0f == 0 = Nothing+ | otherwise = let+ modifier = fromIntegral ((x .&. 0xf0) `shiftR` 4) :: Int16+ (identifier, rem) = if modifier == 0 then readIntFromBytes @Int16 xs else (lastFieldId + modifier, xs)+ elemType = toTType (x .&. 0x0f)+ in Just (rem, elemType, identifier)++readAndAdvance :: IORef Int -> Ptr b -> IO Word8+readAndAdvance bufferPos buffer = do+ pos <- readIORef bufferPos+ b <- peekByteOff buffer pos :: IO Word8+ modifyIORef bufferPos (+ 1)+ return b++readSingleByte :: Int64 -> Ptr b -> IO Word8+readSingleByte pos buffer = peekByteOff buffer (fromIntegral pos)++readNoAdvance :: IORef Int -> Ptr b -> IO Word8+readNoAdvance bufferPos buffer = do+ pos <- readIORef bufferPos+ peekByteOff buffer pos :: IO Word8++compactBooleanTrue :: Word8+compactBooleanTrue = 0x01++compactBooleanFalse :: Word8+compactBooleanFalse = 0x02++compactByte :: Word8+compactByte = 0x03++compactI16 :: Word8+compactI16 = 0x04++compactI32 :: Word8+compactI32 = 0x05++compactI64 :: Word8+compactI64 = 0x06++compactDouble :: Word8+compactDouble = 0x07++compactBinary :: Word8+compactBinary = 0x08++compactList :: Word8+compactList = 0x09++compactSet :: Word8+compactSet = 0x0A++compactMap :: Word8+compactMap = 0x0B++compactStruct :: Word8+compactStruct = 0x0C++compactUuid :: Word8+compactUuid = 0x0D++data TType+ = STOP+ | BOOL+ | BYTE+ | I16+ | I32+ | I64+ | DOUBLE+ | STRING+ | LIST+ | SET+ | MAP+ | STRUCT+ | UUID+ deriving (Show, Eq)++toTType :: Word8 -> TType+toTType t =+ fromMaybe STOP $+ M.lookup (t .&. 0x0f) $+ M.fromList+ [ (compactBooleanTrue, BOOL),+ (compactBooleanFalse, BOOL),+ (compactByte, BYTE),+ (compactI16, I16),+ (compactI32, I32),+ (compactI64, I64),+ (compactDouble, DOUBLE),+ (compactBinary, STRING),+ (compactList, LIST),+ (compactSet, SET),+ (compactMap, MAP),+ (compactStruct, STRUCT),+ (compactUuid, UUID)+ ]++readIntFromBuffer :: (Integral a) => Ptr b -> IORef Int -> IO a+readIntFromBuffer buf bufferPos = do+ n <- readVarIntFromBuffer buf bufferPos+ let u = fromIntegral n :: Word32+ return $ fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1))++readIntFromBytes :: (Integral a) => [Word8] -> (a, [Word8])+readIntFromBytes bs = let+ (n, rem) = readVarIntFromBytes bs+ u = fromIntegral n :: Word32+ in (fromIntegral $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), rem)++readInt32FromBuffer :: Ptr b -> IORef Int -> IO Int32+readInt32FromBuffer buf bufferPos = do+ n <- (fromIntegral <$> readVarIntFromBuffer @Int64 buf bufferPos) :: IO Int32+ let u = fromIntegral n :: Word32+ return $ (fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1))++readInt32FromBytes :: [Word8] -> (Int32, [Word8])+readInt32FromBytes bs = let+ (n', rem) = readVarIntFromBytes @Int64 bs+ n = fromIntegral n' :: Int32+ u = fromIntegral n :: Word32+ in ((fromIntegral (u `shiftR` 1) :: Int32) .^. (-(n .&. 1)), rem)++readVarIntFromBuffer :: (Integral a) => Ptr b -> IORef Int -> IO a+readVarIntFromBuffer buf bufferPos = do+ start <- readIORef bufferPos+ let loop i shift result = do+ b <- readAndAdvance bufferPos buf+ let res = result .|. ((fromIntegral (b .&. 0x7f) :: Integer) `shiftL` shift)+ if (b .&. 0x80) /= 0x80+ then return res+ else loop (i + 1) (shift + 7) res+ fromIntegral <$> loop start 0 0++readVarIntFromBytes :: (Integral a) => [Word8] -> (a, [Word8])+readVarIntFromBytes bs = (fromIntegral n, rem)+ where+ (n, rem) = loop 0 0 bs+ loop _ result [] = (result, [])+ loop shift result (x:xs) = let+ res = result .|. ((fromIntegral (x .&. 0x7f) :: Integer) `shiftL` shift)+ in if (x .&. 0x80) /= 0x80 then (res, xs) else loop (shift + 7) res xs++littleEndianWord8 :: [Word8] -> Word8+littleEndianWord8 bytes+ | length bytes == 4 = foldr (\v acc -> acc .|. v) 0 (zipWith (\b i -> b `shiftL` i) bytes [0,8..])+ | otherwise = error "Expected exactly 4 bytes"++littleEndianWord32 :: [Word8] -> Word32+littleEndianWord32 bytes+ | length bytes == 4 = foldr (\v acc -> acc .|. v) 0 (zipWith (\b i -> (fromIntegral b) `shiftL` i) bytes [0,8..])+ | length bytes < 4 = littleEndianWord32 (take 4 $ bytes ++ (cycle[0]))+ | otherwise = error $ "Expected exactly 4 bytes for Word32 but got " ++ (show bytes)++littleEndianWord64 :: [Word8] -> Word64+littleEndianWord64 bytes+ | length bytes == 8 = foldr (\v acc -> acc .|. v) 0 (zipWith (\b i -> (fromIntegral b) `shiftL` i) bytes [0,8..])+ | otherwise = error "Expected exactly 8 bytes"++littleEndianInt32 :: [Word8] -> Int32+littleEndianInt32 bytes+ | length bytes == 4 = foldr (\v acc -> acc .|. v) 0 (zipWith (\b i -> (fromIntegral b) `shiftL` i) bytes [0,8..])+ | otherwise = error "Expected exactly 4 bytes for Int32"++readUVarInt :: [Word8] -> (Word64, [Word8])+readUVarInt xs = loop xs 0 0 0+ where loop bs x _ 10 = (x, bs)+ loop (b:bs) x s i+ | b < 0x80 = (x .|. ((fromIntegral b) `shiftL` s), bs)+ | otherwise = loop bs (x .|. (fromIntegral ((b .&. 0x7f) `shiftL` s))) (s + 7) (i + 1)++bitStream :: [Word8] -> [[Word8]]+bitStream xs = map (reverse . toBits) xs++toBits :: Word8 -> [Word8]+toBits b = go 1 b+ where+ go 8 n = [(n .&. 1)]+ go i n = (n .&. 1) : go (i + 1) (n .>>. 1)++bitStreamToInt :: Word8 -> [Word8] -> [Int32]+bitStreamToInt _ [] = []+bitStreamToInt bitWidth bits = let+ currBits = take (fromIntegral bitWidth) bits+ remaining = drop (fromIntegral bitWidth) bits+ in bitsToInt32 bitWidth currBits : bitStreamToInt bitWidth remaining++bitsToInt32 :: Word8 -> [Word8] -> Int32+bitsToInt32 bitWidth bits = fromIntegral $ foldr (.|.) 0 (zipWith (\s b -> b .<<. s) [0..] (reverse bits))
src/DataFrame/Internal/Column.hs view
@@ -4,7 +4,7 @@ {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-}-{-# LANGUAGE StrictData #-}+{-# LANGUAGE Strict #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-}@@ -44,8 +44,10 @@ import Control.Exception (throw) import Data.Kind (Type, Constraint) --- | Our representation of a column is a GADT that can store data in either--- a vector with boxed elements or+-- | Our representation of a column is a GADT that can store data based on the underlying data.+-- +-- This allows us to pattern match on data kinds and limit some operations to only some+-- kinds of vectors. E.g. operations for missing data only happen in an OptionalColumn. data Column where BoxedColumn :: Columnable a => VB.Vector a -> Column UnboxedColumn :: (Columnable a, VU.Unbox a) => VU.Vector a -> Column@@ -53,30 +55,45 @@ GroupedBoxedColumn :: Columnable a => VB.Vector (VB.Vector a) -> Column GroupedUnboxedColumn :: (Columnable a, VU.Unbox a) => VB.Vector (VU.Vector a) -> Column GroupedOptionalColumn :: (Columnable a) => VB.Vector (VB.Vector (Maybe a)) -> Column+ -- These are used purely for I/O, not to store live data. MutableBoxedColumn :: Columnable a => VBM.IOVector a -> Column MutableUnboxedColumn :: (Columnable a, VU.Unbox a) => VUM.IOVector a -> Column +-- | A TypedColumn is a wrapper around our type-erased column.+-- It is used to type check expressions on columns. data TypedColumn a where TColumn :: Columnable a => Column -> TypedColumn a +-- | Gets the underlying value from a TypedColumn. unwrapTypedColumn :: TypedColumn a -> Column unwrapTypedColumn (TColumn value) = value --- Functions about column metadata.+-- | An internal function that checks if a column+-- can be used for aggregation. isGrouped :: Column -> Bool isGrouped (GroupedBoxedColumn column) = True isGrouped (GroupedUnboxedColumn column) = True isGrouped _ = False +-- | An internal function that checks if a column can+-- be used in missing value operations.+isOptional :: Column -> Bool+isOptional (OptionalColumn column) = True+isOptional _ = False++-- | An internal/debugging function to get the column type of a column. columnVersionString :: Column -> String columnVersionString column = case column of- BoxedColumn _ -> "Boxed"- UnboxedColumn _ -> "Unboxed"- OptionalColumn _ -> "Optional"- GroupedBoxedColumn _ -> "Grouped Boxed"- GroupedUnboxedColumn _ -> "Grouped Unboxed"+ BoxedColumn _ -> "Boxed"+ UnboxedColumn _ -> "Unboxed"+ OptionalColumn _ -> "Optional"+ GroupedBoxedColumn _ -> "Grouped Boxed"+ GroupedUnboxedColumn _ -> "Grouped Unboxed" GroupedOptionalColumn _ -> "Grouped Optional"+ _ -> "Unknown column string" +-- | An internal/debugging function to get the type stored in the outermost vector+-- of a column. columnTypeString :: Column -> String columnTypeString column = case column of BoxedColumn (column :: VB.Vector a) -> show (typeRep @a)@@ -127,6 +144,8 @@ Just Refl -> VB.map (L.sort . VG.toList) a == VB.map (L.sort . VG.toList) b (==) _ _ = False +-- | A type with column representations used to select the+-- "right" representation when specializing the `toColumn` function. data Rep = RBoxed | RUnboxed@@ -135,10 +154,12 @@ | RGUnboxed | RGOptional +-- | Type-level if statement. type family If (cond :: Bool) (yes :: k) (no :: k) :: k where If 'True yes _ = yes If 'False _ no = no +-- | All unboxable types (according to the `vector` package). type family Unboxable (a :: Type) :: Bool where Unboxable Int = 'True Unboxable Int8 = 'True@@ -163,9 +184,12 @@ KindOf (VU.Vector a) = 'RGUnboxed KindOf a = If (Unboxable a) 'RUnboxed 'RBoxed +-- | A class for converting a vector to a column of the appropriate type.+-- Given each Rep we tell the `toColumnRep` function which Column type to pick. class ColumnifyRep (r :: Rep) a where toColumnRep :: VB.Vector a -> Column +-- | Constraint synonym for what we can put into columns. type Columnable a = (Columnable' a, ColumnifyRep (KindOf a) a, UnboxIf a, SBoolI (Unboxable a) ) instance (Columnable a, VU.Unbox a)@@ -188,16 +212,49 @@ => ColumnifyRep 'RGUnboxed (VU.Vector a) where toColumnRep = GroupedUnboxedColumn -toColumn' ::+{- | O(n) Convert a vector to a column. Automatically picks the best representation of a vector to store the underlying data in.++__Examples:__++@+> import qualified Data.Vector as V+> fromVector (V.fromList [(1 :: Int), 2, 3, 4])+[1,2,3,4]+@+-}+fromVector :: forall a. (Columnable a, ColumnifyRep (KindOf a) a) => VB.Vector a -> Column-toColumn' = toColumnRep @(KindOf a)+fromVector = toColumnRep @(KindOf a) -toColumn ::+{- | O(n) Convert an unboxed vector to a column. This avoids the extra conversion if you already have the data in an unboxed vector.++__Examples:__++@+> import qualified Data.Vector.Unboxed as V+> fromVector (V.fromList [(1 :: Int), 2, 3, 4])+[1,2,3,4]+@+-}+fromUnboxedVector :: forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column+fromUnboxedVector = UnboxedColumn++{- | O(n) Convert a list to a column. Automatically picks the best representation of a vector to store the underlying data in.++__Examples:__++@+> fromList [(1 :: Int), 2, 3, 4]+[1,2,3,4]+@+-}+fromList :: forall a. (Columnable a, ColumnifyRep (KindOf a) a) => [a] -> Column-toColumn = toColumnRep @(KindOf a) . VB.fromList+fromList = toColumnRep @(KindOf a) . VB.fromList + data SBool (b :: Bool) where STrue :: SBool 'True SFalse :: SBool 'False@@ -217,52 +274,44 @@ type UnboxIf a = When (Unboxable a) (VU.Unbox a) --- | Generic column transformation (no index).-transform+-- | An internal function to map a function over the values of a column.+mapColumn :: forall b c. ( Columnable b , Columnable c- , UnboxIf c- , Typeable b- , Typeable c )+ , UnboxIf c) => (b -> c) -> Column -> Maybe Column-transform f = \case+mapColumn f = \case BoxedColumn (col :: VB.Vector a) | Just Refl <- testEquality (typeRep @a) (typeRep @b)- -> Just (toColumn' @c (VB.map f col))+ -> Just (fromVector @c (VB.map f col)) | otherwise -> Nothing OptionalColumn (col :: VB.Vector a) | Just Refl <- testEquality (typeRep @a) (typeRep @b)- -> Just (toColumn' @c (VB.map f col))+ -> Just (fromVector @c (VB.map f col)) | otherwise -> Nothing UnboxedColumn (col :: VU.Vector a) | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> Just $ case sUnbox @c of STrue -> UnboxedColumn (VU.map f col)- SFalse -> toColumn' @c (VB.map f (VB.convert col))+ SFalse -> fromVector @c (VB.map f (VB.convert col)) | otherwise -> Nothing GroupedBoxedColumn (col :: VB.Vector (VB.Vector a)) | Just Refl <- testEquality (typeRep @(VB.Vector a)) (typeRep @b)- -> Just (toColumn' @c (VB.map f col))+ -> Just (fromVector @c (VB.map f col)) | otherwise -> Nothing GroupedUnboxedColumn (col :: VB.Vector (VU.Vector a)) | Just Refl <- testEquality (typeRep @(VU.Vector a)) (typeRep @b)- -> Just (toColumn' @c (VB.map f col))+ -> Just (fromVector @c (VB.map f col)) | otherwise -> Nothing GroupedOptionalColumn (col :: VB.Vector (VB.Vector a)) | Just Refl <- testEquality (typeRep @(VB.Vector a)) (typeRep @b)- -> Just (toColumn' @c (VB.map f col))+ -> Just (fromVector @c (VB.map f col)) | otherwise -> Nothing --- | Converts a an unboxed vector to a column making sure to put--- the vector into an appropriate column type by reflection on the--- vector's type parameter.-toColumnUnboxed :: forall a. (Columnable a, VU.Unbox a) => VU.Vector a -> Column-toColumnUnboxed = UnboxedColumn --- Functions that don't depend on column type. -- | O(1) Gets the number of elements in the column. columnLength :: Column -> Int columnLength (BoxedColumn xs) = VG.length xs@@ -273,6 +322,7 @@ columnLength (GroupedOptionalColumn xs) = VG.length xs {-# INLINE columnLength #-} +-- | O(n) Takes the first n values of a column. takeColumn :: Int -> Column -> Column takeColumn n (BoxedColumn xs) = BoxedColumn $ VG.take n xs takeColumn n (UnboxedColumn xs) = UnboxedColumn $ VG.take n xs@@ -282,12 +332,12 @@ takeColumn n (GroupedOptionalColumn xs) = GroupedOptionalColumn $ VG.take n xs {-# INLINE takeColumn #-} --- TODO: Maybe we can remvoe all this boilerplate and make--- transform take in a generic vector function.+-- | O(n) Takes the last n values of a column. takeLastColumn :: Int -> Column -> Column takeLastColumn n column = sliceColumn (columnLength column - n) n column {-# INLINE takeLastColumn #-} +-- | O(n) Takes n values after a given column index. sliceColumn :: Int -> Int -> Column -> Column sliceColumn start n (BoxedColumn xs) = BoxedColumn $ VG.slice start n xs sliceColumn start n (UnboxedColumn xs) = UnboxedColumn $ VG.slice start n xs@@ -297,7 +347,7 @@ sliceColumn start n (GroupedOptionalColumn xs) = GroupedOptionalColumn $ VG.slice start n xs {-# INLINE sliceColumn #-} --- TODO: We can probably generalize this to `applyVectorFunction`.+-- | O(n) Selects the elements at a given set of indices. May change the order. atIndices :: S.Set Int -> Column -> Column atIndices indexes (BoxedColumn column) = BoxedColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column atIndices indexes (OptionalColumn column) = OptionalColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column@@ -307,6 +357,7 @@ atIndices indexes (GroupedOptionalColumn column) = GroupedOptionalColumn $ VG.ifilter (\i _ -> i `S.member` indexes) column {-# INLINE atIndices #-} +-- | O(n) Selects the elements at a given set of indices. Does not change the order. atIndicesStable :: VU.Vector Int -> Column -> Column atIndicesStable indexes (BoxedColumn column) = BoxedColumn $ indexes `getIndices` column atIndicesStable indexes (UnboxedColumn column) = UnboxedColumn $ indexes `getIndicesUnboxed` column@@ -316,14 +367,17 @@ atIndicesStable indexes (GroupedOptionalColumn column) = GroupedOptionalColumn $ indexes `getIndices` column {-# INLINE atIndicesStable #-} +-- | Internal helper to get indices in a boxed vector. getIndices :: VU.Vector Int -> VB.Vector a -> VB.Vector a getIndices indices xs = VB.generate (VU.length indices) (\i -> xs VB.! (indices VU.! i)) {-# INLINE getIndices #-} +-- | Internal helper to get indices in an unboxed vector. getIndicesUnboxed :: (VU.Unbox a) => VU.Vector Int -> VU.Vector a -> VU.Vector a getIndicesUnboxed indices xs = VU.generate (VU.length indices) (\i -> xs VU.! (indices VU.! i)) {-# INLINE getIndicesUnboxed #-} +-- | An internal function that returns a vector of how indexes change after a column is sorted. sortedIndexes :: Bool -> Column -> VU.Vector Int sortedIndexes asc (BoxedColumn column ) = runST $ do withIndexes <- VG.thaw $ VG.indexed column@@ -358,32 +412,32 @@ {-# INLINE sortedIndexes #-} -- | Applies a function that returns an unboxed result to an unboxed vector, storing the result in a column.-itransform- :: forall b c. (Typeable b, Typeable c, Columnable b, Columnable c)+imapColumn+ :: forall b c. (Columnable b, Columnable c) => (Int -> b -> c) -> Column -> Maybe Column-itransform f = \case+imapColumn f = \case BoxedColumn (col :: VB.Vector a) | Just Refl <- testEquality (typeRep @a) (typeRep @b)- -> Just (toColumn' @c (VB.imap f col))+ -> Just (fromVector @c (VB.imap f col)) | otherwise -> Nothing UnboxedColumn (col :: VU.Vector a) | Just Refl <- testEquality (typeRep @a) (typeRep @b) -> Just $ case sUnbox @c of STrue -> UnboxedColumn (VU.imap f col)- SFalse -> toColumn' @c (VB.imap f (VB.convert col))+ SFalse -> fromVector @c (VB.imap f (VB.convert col)) | otherwise -> Nothing OptionalColumn (col :: VB.Vector a) | Just Refl <- testEquality (typeRep @a) (typeRep @b)- -> Just (toColumn' @c (VB.imap f col))+ -> Just (fromVector @c (VB.imap f col)) | otherwise -> Nothing GroupedBoxedColumn (col :: VB.Vector a) | Just Refl <- testEquality (typeRep @a) (typeRep @b)- -> Just (toColumn' @c (VB.imap f col))+ -> Just (fromVector @c (VB.imap f col)) | otherwise -> Nothing GroupedUnboxedColumn (col :: VB.Vector a) | Just Refl <- testEquality (typeRep @a) (typeRep @b)- -> Just (toColumn' @c (VB.imap f col))+ -> Just (fromVector @c (VB.imap f col)) | otherwise -> Nothing -- | Filter column with index.@@ -401,7 +455,7 @@ Refl <- testEquality (typeRep @a) (typeRep @b) return $ GroupedUnboxedColumn $ VG.ifilter f column -+-- | Fold (right) column with index. ifoldrColumn :: forall a b. (Columnable a, Columnable b) => (Int -> a -> b -> b) -> b -> Column -> Maybe b ifoldrColumn f acc c@(BoxedColumn (column :: VB.Vector d)) = do Refl <- testEquality (typeRep @a) (typeRep @d)@@ -419,6 +473,7 @@ Refl <- testEquality (typeRep @a) (typeRep @d) return $ VG.ifoldr f acc column +-- | Fold (left) column with index. ifoldlColumn :: forall a b . (Columnable a, Columnable b) => (b -> Int -> a -> b) -> b -> Column -> Maybe b ifoldlColumn f acc c@(BoxedColumn (column :: VB.Vector d)) = do Refl <- testEquality (typeRep @a) (typeRep @d)@@ -436,33 +491,24 @@ Refl <- testEquality (typeRep @a) (typeRep @d) return $ VG.ifoldl' f acc column -reduceColumn :: forall a b. Columnable a => (a -> b) -> Column -> b+-- | Generic reduce function for all Column types.+reduceColumn :: forall a b. Columnable a => (a -> b) -> Column -> Maybe b {-# SPECIALIZE reduceColumn ::- (VU.Vector (Double, Double) -> Double) -> Column -> Double,- (VU.Vector Double -> Double) -> Column -> Double #-}-reduceColumn f (BoxedColumn (column :: c)) = case testEquality (typeRep @c) (typeRep @a) of- Just Refl -> f column- Nothing -> error $ "Can't reduce. Incompatible types: " ++ show (typeRep @a) ++ " " ++ show (typeRep @a)-reduceColumn f (UnboxedColumn (column :: c)) = case testEquality (typeRep @c) (typeRep @a) of- Just Refl -> f column- Nothing -> error $ "Can't reduce. Incompatible types: " ++ show (typeRep @a) ++ " " ++ show (typeRep @a)-reduceColumn f (OptionalColumn (column :: c)) = case testEquality (typeRep @c) (typeRep @a) of- Just Refl -> f column- Nothing -> error $ "Can't reduce. Incompatible types: " ++ show (typeRep @a) ++ " " ++ show (typeRep @a)-{-# INLINE reduceColumn #-}--safeReduceColumn :: forall a b. (Typeable a) => (a -> b) -> Column -> Maybe b-safeReduceColumn f (BoxedColumn (column :: c)) = do+ (VU.Vector (Double, Double) -> Double) -> Column -> Maybe Double,+ (VU.Vector Double -> Double) -> Column -> Maybe Double #-}+reduceColumn f (BoxedColumn (column :: c)) = do Refl <- testEquality (typeRep @c) (typeRep @a)- return $! f column-safeReduceColumn f (UnboxedColumn (column :: c)) = do+ pure $ f column+reduceColumn f (UnboxedColumn (column :: c)) = do Refl <- testEquality (typeRep @c) (typeRep @a)- return $! f column-safeReduceColumn f (OptionalColumn (column :: c)) = do+ pure $ f column+reduceColumn f (OptionalColumn (column :: c)) = do Refl <- testEquality (typeRep @c) (typeRep @a)- return $! f column-{-# INLINE safeReduceColumn #-}+ pure $ f column+reduceColumn _ _ = error "UNIMPLEMENTED"+{-# INLINE reduceColumn #-} +-- | An internal, column version of zip. zipColumns :: Column -> Column -> Column zipColumns (BoxedColumn column) (BoxedColumn other) = BoxedColumn (VG.zip column other) zipColumns (BoxedColumn column) (UnboxedColumn other) = BoxedColumn (VB.generate (min (VG.length column) (VG.length other)) (\i -> (column VG.! i, other VG.! i)))@@ -470,20 +516,22 @@ zipColumns (UnboxedColumn column) (UnboxedColumn other) = UnboxedColumn (VG.zip column other) {-# INLINE zipColumns #-} -zipWithColumns :: forall a b c . (Columnable a, Columnable b, Columnable c) => (a -> b -> c) -> Column -> Column -> Column+-- | An internal, column version of zipWith.+zipWithColumns :: forall a b c . (Columnable a, Columnable b, Columnable c) => (a -> b -> c) -> Column -> Column -> Maybe Column zipWithColumns f (UnboxedColumn (column :: VU.Vector d)) (UnboxedColumn (other :: VU.Vector e)) = case testEquality (typeRep @a) (typeRep @d) of Just Refl -> case testEquality (typeRep @b) (typeRep @e) of- Just Refl -> toColumn' $ VB.zipWith f (VG.convert column) (VG.convert other)- Nothing -> throw $ TypeMismatchException' (typeRep @b) (show $ typeRep @e) "" "zipWithColumns"- Nothing -> throw $ TypeMismatchException' (typeRep @a) (show $ typeRep @d) "" "zipWithColumns"+ Just Refl -> pure $ case sUnbox @c of+ STrue -> fromUnboxedVector (VU.zipWith f column other)+ SFalse -> fromVector $ VB.zipWith f (VG.convert column) (VG.convert other)+ Nothing -> Nothing+ Nothing -> Nothing zipWithColumns f left right = let left' = toVector @a left right' = toVector @b right- in toColumn' $ VB.zipWith f left' right' + in pure $ fromVector $ VB.zipWith f left' right' {-# INLINE zipWithColumns #-} -- Functions for mutable columns (intended for IO).--- Clean this up. writeColumn :: Int -> T.Text -> Column -> IO (Either T.Text Bool) writeColumn i value (MutableBoxedColumn (col :: VBM.IOVector a)) = let in case testEquality (typeRep @a) (typeRep @T.Text) of@@ -514,31 +562,97 @@ | otherwise = VU.unsafeFreeze col >>= \c -> return $ BoxedColumn $ VB.generate (VU.length c) (\i -> if i `elem` map fst nulls then Left (fromMaybe (error "") (lookup i nulls)) else Right (c VU.! i)) {-# INLINE freezeColumn' #-} +-- | Fills the end of a column, up to n, with Nothing. Does nothing if column has length greater than n. expandColumn :: Int -> Column -> Column expandColumn n (OptionalColumn col) = OptionalColumn $ col <> VB.replicate n Nothing-expandColumn n (BoxedColumn col) = OptionalColumn $ VB.map Just col <> VB.replicate n Nothing-expandColumn n (UnboxedColumn col) = OptionalColumn $ VB.map Just (VU.convert col) <> VB.replicate n Nothing-expandColumn n (GroupedBoxedColumn col) = GroupedBoxedColumn $ col <> VB.replicate n VB.empty-expandColumn n (GroupedUnboxedColumn col) = GroupedUnboxedColumn $ col <> VB.replicate n VU.empty+expandColumn n column@(BoxedColumn col)+ | n < VG.length col = OptionalColumn $ VB.map Just col <> VB.replicate n Nothing+ | otherwise = column+expandColumn n column@(UnboxedColumn col)+ | n < VG.length col = OptionalColumn $ VB.map Just (VU.convert col) <> VB.replicate n Nothing+ | otherwise = column+expandColumn n column@(GroupedBoxedColumn col)+ | n < VG.length col = GroupedBoxedColumn $ col <> VB.replicate n VB.empty+ | otherwise = column+expandColumn n column@(GroupedUnboxedColumn col)+ | n < VG.length col = GroupedUnboxedColumn $ col <> VB.replicate n VU.empty+ | otherwise = column +-- | Fills the beginning of a column, up to n, with Nothing. Does nothing if column has length greater than n.+leftExpandColumn :: Int -> Column -> Column+leftExpandColumn n (OptionalColumn col) = OptionalColumn $ VB.replicate n Nothing <> col+leftExpandColumn n (BoxedColumn col) = OptionalColumn $ VB.replicate n Nothing <> VB.map Just col+leftExpandColumn n (UnboxedColumn col) = OptionalColumn $ VB.replicate n Nothing <> VB.map Just (VU.convert col)+leftExpandColumn n (GroupedBoxedColumn col) = GroupedBoxedColumn $ VB.replicate n VB.empty <> col+leftExpandColumn n (GroupedUnboxedColumn col) = GroupedUnboxedColumn $ VB.replicate n VU.empty <> col++-- | Concatenates two columns.+concatColumns :: Column -> Column -> Maybe Column+concatColumns (OptionalColumn left) (OptionalColumn right) = case testEquality (typeOf left) (typeOf right) of+ Nothing -> Nothing+ Just Refl -> Just (OptionalColumn $ left <> right)+concatColumns (BoxedColumn left) (BoxedColumn right) = case testEquality (typeOf left) (typeOf right) of+ Nothing -> Nothing+ Just Refl -> Just (BoxedColumn $ left <> right)+concatColumns (UnboxedColumn left) (UnboxedColumn right) = case testEquality (typeOf left) (typeOf right) of+ Nothing -> Nothing+ Just Refl -> Just (UnboxedColumn $ left <> right)+concatColumns (GroupedBoxedColumn left) (GroupedBoxedColumn right) = case testEquality (typeOf left) (typeOf right) of+ Nothing -> Nothing+ Just Refl -> Just (GroupedBoxedColumn $ left <> right)+concatColumns (GroupedUnboxedColumn left) (GroupedUnboxedColumn right) = case testEquality (typeOf left) (typeOf right) of+ Nothing -> Nothing+ Just Refl -> Just (GroupedUnboxedColumn $ left <> right)+concatColumns _ _ = Nothing++{- | O(n) Converts a column to a boxed vector. Throws an exception if the wrong type is specified.++__Examples:__++@+> column = fromList [(1 :: Int), 2, 3, 4]+> toVector @Int column+[1,2,3,4]+> toVector @Double column+exception: ...+-} toVector :: forall a . Columnable a => Column -> VB.Vector a-toVector column@(OptionalColumn (col :: VB.Vector b)) =+toVector = toVectorWithLabel "toVector"++-- | An internal version of toVector that takes the calling function as an extra argument.+toVectorWithLabel :: forall a . Columnable a => String -> Column -> VB.Vector a+toVectorWithLabel label column@(OptionalColumn (col :: VB.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of Just Refl -> col- Nothing -> throw $ TypeMismatchException' (typeRep @b) (show $ typeRep @a) "" "toVector"-toVector (BoxedColumn (col :: VB.Vector b)) =+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a)+ , expectedType = Right (typeRep @b)+ , callingFunctionName = Just label+ , errorColumnName = Nothing})+toVectorWithLabel label (BoxedColumn (col :: VB.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of Just Refl -> col- Nothing -> throw $ TypeMismatchException' (typeRep @b) (show $ typeRep @a) "" "toVector"-toVector (UnboxedColumn (col :: VU.Vector b)) =+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a)+ , expectedType = Right (typeRep @b)+ , callingFunctionName = Just label+ , errorColumnName = Nothing})+toVectorWithLabel label (UnboxedColumn (col :: VU.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of Just Refl -> VB.convert col- Nothing -> throw $ TypeMismatchException' (typeRep @b) (show $ typeRep @a) "" "toVector"-toVector (GroupedBoxedColumn (col :: VB.Vector b)) =+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a)+ , expectedType = Right (typeRep @b)+ , callingFunctionName = Just label+ , errorColumnName = Nothing})+toVectorWithLabel label (GroupedBoxedColumn (col :: VB.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of Just Refl -> col- Nothing -> throw $ TypeMismatchException' (typeRep @b) (show $ typeRep @a) "" "toVector"-toVector (GroupedUnboxedColumn (col :: VB.Vector b)) =+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a)+ , expectedType = Right (typeRep @b)+ , callingFunctionName = Just label+ , errorColumnName = Nothing})+toVectorWithLabel label (GroupedUnboxedColumn (col :: VB.Vector b)) = case testEquality (typeRep @a) (typeRep @b) of Just Refl -> col- Nothing -> throw $ TypeMismatchException' (typeRep @b) (show $ typeRep @a) "" "toVector"+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext { userType = Right (typeRep @a)+ , expectedType = Right (typeRep @b)+ , callingFunctionName = Just label+ , errorColumnName = Nothing})
src/DataFrame/Internal/DataFrame.hs view
@@ -4,7 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE GADTs #-}-{-# LANGUAGE StrictData #-}+{-# LANGUAGE Strict #-} {-# LANGUAGE FlexibleContexts #-} module DataFrame.Internal.DataFrame where
src/DataFrame/Internal/Expression.hs view
@@ -32,18 +32,19 @@ Apply :: (Columnable a, Columnable b) => T.Text -> (b -> a) -> Expr b -> Expr a BinOp :: (Columnable c, Columnable b, Columnable a) => T.Text -> (c -> b -> a) -> Expr c -> Expr b -> Expr a -interpret :: forall a b . (Columnable a) => DataFrame -> Expr a -> TypedColumn a-interpret df (Lit value) = TColumn $ toColumn' $ V.replicate (fst $ dataframeDimensions df) value+interpret :: forall a . (Columnable a) => DataFrame -> Expr a -> TypedColumn a+interpret df (Lit value) = TColumn $ fromVector $ V.replicate (fst $ dataframeDimensions df) value interpret df (Col name) = case getColumn name df of Nothing -> throw $ ColumnNotFoundException name "" (map fst $ M.toList $ columnIndices df) Just col -> TColumn col interpret df (Apply _ (f :: c -> d) value) = let (TColumn value') = interpret @c df value- in TColumn $ fromMaybe (error "transform returned nothing") (transform f value')+ -- TODO: Handle this gracefully.+ in TColumn $ fromMaybe (error "mapColumn returned nothing") (mapColumn f value') interpret df (BinOp _ (f :: c -> d -> e) left right) = let (TColumn left') = interpret @c df left (TColumn right') = interpret @d df right- in TColumn $ zipWithColumns f left' right'+ in TColumn $ fromMaybe (error "mapColumn returned nothing") (zipWithColumns f left' right') instance (Num a, Columnable a) => Num (Expr a) where (+) :: Expr a -> Expr a -> Expr a
src/DataFrame/Internal/Types.hs view
@@ -19,4 +19,4 @@ import Type.Reflection (TypeRep, typeOf, typeRep) import Data.Type.Equality (TestEquality(..)) -type Columnable' a = (Typeable a, Show a, Ord a, Eq a)+type Columnable' a = (Typeable a, Show a, Ord a, Eq a, Read a)
+ src/DataFrame/Lazy.hs view
@@ -0,0 +1,3 @@+module DataFrame.Lazy (module DataFrame.Lazy.Internal.DataFrame) where++import DataFrame.Lazy.Internal.DataFrame
+ src/DataFrame/Lazy/IO/CSV.hs view
@@ -0,0 +1,327 @@+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE ExplicitNamespaces #-}+{-# LANGUAGE LambdaCase #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE ScopedTypeVariables #-}+{-# LANGUAGE TypeApplications #-}+{-# LANGUAGE GADTs #-}+{-# LANGUAGE RankNTypes #-}+{-# LANGUAGE Strict #-}+module DataFrame.Lazy.IO.CSV where++import qualified Data.ByteString.Char8 as C+import qualified Data.List as L+import qualified Data.Map as M+import qualified Data.Set as S+import qualified Data.Text as T+import qualified Data.Text.Lazy as TL+import qualified Data.Text.Lazy.IO as TLIO+import qualified Data.Text.IO as TIO+import qualified Data.Vector as V+import qualified Data.Vector.Unboxed as VU+import qualified Data.Vector.Mutable as VM+import qualified Data.Vector.Unboxed.Mutable as VUM++import Control.Applicative ((<$>), (<|>), (<*>), (<*), (*>), many)+import Control.Monad (forM_, zipWithM_, unless, when, void, replicateM_)+import Data.Attoparsec.Text+import Data.Char+import DataFrame.Internal.Column (Column(..), freezeColumn', writeColumn, columnLength)+import DataFrame.Internal.DataFrame (DataFrame(..))+import DataFrame.Internal.Parsing+import DataFrame.Operations.Typing+import Data.Foldable (fold)+import Data.Function (on)+import Data.IORef+import Data.Maybe+import Data.Text.Encoding (decodeUtf8Lenient)+import Data.Type.Equality+ ( TestEquality (testEquality),+ type (:~:) (Refl)+ )+import GHC.IO.Handle (Handle)+import Prelude hiding (concat, takeWhile)+import System.IO+import Type.Reflection++-- | Record for CSV read options.+data ReadOptions = ReadOptions {+ hasHeader :: Bool,+ inferTypes :: Bool,+ safeRead :: Bool,+ rowRange :: !(Maybe (Int, Int)), -- (start, length)+ seekPos :: !(Maybe Integer),+ totalRows :: !(Maybe Int),+ leftOver :: !T.Text,+ rowsRead :: !Int+}++-- | By default we assume the file has a header, we infer the types on read+-- and we convert any rows with nullish objects into Maybe (safeRead).+defaultOptions :: ReadOptions+defaultOptions = ReadOptions { hasHeader = True, inferTypes = True, safeRead = True, rowRange = Nothing, seekPos = Nothing, totalRows = Nothing, leftOver = "", rowsRead = 0 }++-- | Reads a CSV file from the given path.+-- Note this file stores intermediate temporary files+-- while converting the CSV from a row to a columnar format.+readCsv :: String -> IO DataFrame+readCsv path = fst <$> readSeparated ',' defaultOptions path++-- | Reads a tab separated file from the given path.+-- Note this file stores intermediate temporary files+-- while converting the CSV from a row to a columnar format.+readTsv :: String -> IO DataFrame+readTsv path = fst <$> readSeparated '\t' defaultOptions path++-- | Reads a character separated file into a dataframe using mutable vectors.+readSeparated :: Char -> ReadOptions -> String -> IO (DataFrame, (Integer, T.Text, Int))+readSeparated c opts path = do+ totalRows <- case totalRows opts of+ Nothing -> countRows c path >>= \total -> if hasHeader opts then return (total - 1) else return total+ Just n -> if hasHeader opts then return (n - 1) else return n+ let (begin, len) = case rowRange opts of+ Nothing -> (0, totalRows)+ Just (start, len) -> (start, min len (totalRows - rowsRead opts))+ withFile path ReadMode $ \handle -> do+ firstRow <- map T.strip . parseSep c <$> TIO.hGetLine handle+ let columnNames = if hasHeader opts+ then map (T.filter (/= '\"')) firstRow+ else map (T.singleton . intToDigit) [0..(length firstRow - 1)]+ -- If there was no header rewind the file cursor.+ unless (hasHeader opts) $ hSeek handle AbsoluteSeek 0++ currPos <- hTell handle+ when (isJust $ seekPos opts) $ hSeek handle AbsoluteSeek (fromMaybe currPos (seekPos opts))++ -- Initialize mutable vectors for each column+ let numColumns = length columnNames+ let numRows = len+ -- Use this row to infer the types of the rest of the column.+ -- TODO: this isn't robust but in so far as this is a guess anyway+ -- it's probably fine. But we should probably sample n rows and pick+ -- the most likely type from the sample.+ -- dataRow <- map T.strip . parseSep c . (<>) (leftOver opts) <$> TIO.hGetLine handle+ (!dataRow, !remainder) <- readSingleLine c (leftOver opts) handle++ -- This array will track the indices of all null values for each column.+ -- If any exist then the column will be an optional type.+ nullIndices <- VM.unsafeNew numColumns+ VM.set nullIndices []+ mutableCols <- VM.unsafeNew numColumns+ getInitialDataVectors numRows mutableCols dataRow++ -- Read rows into the mutable vectors+ (!unconsumed, !r) <- fillColumns numRows c mutableCols nullIndices remainder handle++ -- Freeze the mutable vectors into immutable ones+ nulls' <- V.unsafeFreeze nullIndices+ cols <- V.mapM (freezeColumn mutableCols nulls' opts) (V.generate numColumns id)+ pos <- hTell handle++ return (DataFrame {+ columns = cols,+ freeIndices = [],+ columnIndices = M.fromList (zip columnNames [0..]),+ dataframeDimensions = (maybe 0 columnLength (cols V.! 0), V.length cols)+ }, (pos, unconsumed, r + 1))+{-# INLINE readSeparated #-}++getInitialDataVectors :: Int -> VM.IOVector Column -> [T.Text] -> IO ()+getInitialDataVectors n mCol xs = do+ forM_ (zip [0..] xs) $ \(i, x) -> do+ col <- case inferValueType x of+ "Int" -> MutableUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Int)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readInt x) >> return c)+ "Double" -> MutableUnboxedColumn <$> ((VUM.unsafeNew n :: IO (VUM.IOVector Double)) >>= \c -> VUM.unsafeWrite c 0 (fromMaybe 0 $ readDouble x) >> return c)+ _ -> MutableBoxedColumn <$> ((VM.unsafeNew n :: IO (VM.IOVector T.Text)) >>= \c -> VM.unsafeWrite c 0 x >> return c)+ VM.unsafeWrite mCol i col+{-# INLINE getInitialDataVectors #-}++inferValueType :: T.Text -> T.Text+inferValueType s = let+ example = s+ in case readInt example of+ Just _ -> "Int"+ Nothing -> case readDouble example of+ Just _ -> "Double"+ Nothing -> "Other"+{-# INLINE inferValueType #-}++readSingleLine :: Char -> T.Text -> Handle -> IO ([T.Text], T.Text)+readSingleLine c unused handle = parseWith (TIO.hGetChunk handle) (parseRow c) unused >>= \case+ Fail unconsumed ctx er -> do+ erpos <- hTell handle+ fail $ "Failed to parse CSV file around " <> show erpos <> " byte; due: "+ <> show er <> "; context: " <> show ctx+ Partial c -> do+ fail "Partial handler is called"+ Done (unconsumed :: T.Text) (row :: [T.Text]) -> do+ return (row, unconsumed)++-- | Reads rows from the handle and stores values in mutable vectors.+fillColumns :: Int -> Char -> VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> T.Text -> Handle -> IO (T.Text, Int)+fillColumns n c mutableCols nullIndices unused handle = do+ input <- newIORef unused+ rowsRead <- newIORef (0 :: Int)+ forM_ [1..(n - 1)] $ \i -> do+ isEOF <- hIsEOF handle+ input' <- readIORef input+ unless (isEOF && input' == mempty) $ do+ parseWith (TIO.hGetChunk handle) (parseRow c) input' >>= \case+ Fail unconsumed ctx er -> do+ erpos <- hTell handle+ fail $ "Failed to parse CSV file around " <> show erpos <> " byte; due: "+ <> show er <> "; context: " <> show ctx+ Partial c -> do+ fail "Partial handler is called"+ Done (unconsumed :: T.Text) (row :: [T.Text]) -> do+ writeIORef input unconsumed+ modifyIORef rowsRead (+1)+ zipWithM_ (writeValue mutableCols nullIndices i) [0..] row+ l <- readIORef input+ r <- readIORef rowsRead+ pure (l, r)+{-# INLINE fillColumns #-}++-- | Writes a value into the appropriate column, resizing the vector if necessary.+writeValue :: VM.IOVector Column -> VM.IOVector [(Int, T.Text)] -> Int -> Int -> T.Text -> IO ()+writeValue mutableCols nullIndices count colIndex value = do+ col <- VM.unsafeRead mutableCols colIndex+ res <- writeColumn count value col+ let modify value = VM.unsafeModify nullIndices ((count, value) :) colIndex+ either modify (const (return ())) res+{-# INLINE writeValue #-}++-- | Freezes a mutable vector into an immutable one, trimming it to the actual row count.+freezeColumn :: VM.IOVector Column -> V.Vector [(Int, T.Text)] -> ReadOptions -> Int -> IO (Maybe Column)+freezeColumn mutableCols nulls opts colIndex = do+ col <- VM.unsafeRead mutableCols colIndex+ Just <$> freezeColumn' (nulls V.! colIndex) col+{-# INLINE freezeColumn #-}++parseSep :: Char -> T.Text -> [T.Text]+parseSep c s = either error id (parseOnly (record c) s)+{-# INLINE parseSep #-}++record :: Char -> Parser [T.Text]+record c =+ field c `sepBy1` char c+ <?> "record"+{-# INLINE record #-}++parseRow :: Char -> Parser [T.Text]+parseRow c = (record c <* lineEnd) <?> "record-new-line"++field :: Char -> Parser T.Text+field c =+ quotedField <|> unquotedField c+ <?> "field"+{-# INLINE field #-}++unquotedTerminators :: Char -> S.Set Char+unquotedTerminators sep = S.fromList [sep, '\n', '\r', '"']++unquotedField :: Char -> Parser T.Text+unquotedField sep =+ takeWhile (not . (`S.member` terminators)) <?> "unquoted field"+ where terminators = unquotedTerminators sep+{-# INLINE unquotedField #-}++quotedField :: Parser T.Text+quotedField = char '"' *> contents <* char '"' <?> "quoted field"+ where+ contents = fold <$> many (unquote <|> unescape)+ where+ unquote = takeWhile1 (notInClass "\"\\")+ unescape = char '\\' *> do+ T.singleton <$> do+ char '\\' <|> char '"'+{-# INLINE quotedField #-}++lineEnd :: Parser ()+lineEnd =+ (endOfLine <|> endOfInput)+ <?> "end of line"+{-# INLINE lineEnd #-}++-- | First pass to count rows for exact allocation+countRows :: Char -> FilePath -> IO Int+countRows c path = withFile path ReadMode $! go 0 ""+ where+ go !n !input h = do+ isEOF <- hIsEOF h+ if isEOF && input == mempty+ then pure n+ else+ parseWith (TIO.hGetChunk h) (parseRow c) input >>= \case+ Fail unconsumed ctx er -> do+ erpos <- hTell h+ fail $ "Failed to parse CSV file around " <> show erpos <> " byte; due: "+ <> show er <> "; context: " <> show ctx <> " " <> show unconsumed+ Partial c -> do+ fail $ "Partial handler is called; n = " <> show n+ Done (unconsumed :: T.Text) _ ->+ go (n + 1) unconsumed h+{-# INLINE countRows #-}++writeCsv :: String -> DataFrame -> IO ()+writeCsv = writeSeparated ','++writeSeparated :: Char -- ^ Separator+ -> String -- ^ Path to write to+ -> DataFrame+ -> IO ()+writeSeparated c filepath df = withFile filepath WriteMode $ \handle ->do+ let (rows, columns) = dataframeDimensions df+ let headers = map fst (L.sortBy (compare `on` snd) (M.toList (columnIndices df)))+ TIO.hPutStrLn handle (T.intercalate ", " headers)+ forM_ [0..(rows - 1)] $ \i -> do+ let row = getRowAsText df i+ TIO.hPutStrLn handle (T.intercalate ", " row)++getRowAsText :: DataFrame -> Int -> [T.Text]+getRowAsText df i = V.ifoldr go [] (columns df)+ where+ indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df))+ go k Nothing acc = acc+ go k (Just (BoxedColumn (c :: V.Vector a))) acc = case c V.!? i of+ Just e -> textRep : acc+ where textRep = case testEquality (typeRep @a) (typeRep @T.Text) of+ Just Refl -> e+ Nothing -> case typeRep @a of+ App t1 t2 -> case eqTypeRep t1 (typeRep @Maybe) of+ Just HRefl -> case testEquality t2 (typeRep @T.Text) of+ Just Refl -> fromMaybe "null" e+ Nothing -> (fromOptional . (T.pack . show)) e+ where fromOptional s+ | T.isPrefixOf "Just " s = T.drop (T.length "Just ") s+ | otherwise = "null"+ Nothing -> (T.pack . show) e+ _ -> (T.pack . show) e+ Nothing ->+ error $+ "Column "+ ++ T.unpack (indexMap M.! k)+ ++ " has less items than "+ ++ "the other columns at index "+ ++ show i+ go k (Just (UnboxedColumn c)) acc = case c VU.!? i of+ Just e -> T.pack (show e) : acc+ Nothing ->+ error $+ "Column "+ ++ T.unpack (indexMap M.! k)+ ++ " has less items than "+ ++ "the other columns at index "+ ++ show i+ go k (Just (OptionalColumn (c :: V.Vector (Maybe a)))) acc = case c V.!? i of+ Just e -> textRep : acc+ where textRep = case testEquality (typeRep @a) (typeRep @T.Text) of+ Just Refl -> fromMaybe "Nothing" e+ Nothing -> (T.pack . show) e+ Nothing ->+ error $+ "Column "+ ++ T.unpack (indexMap M.! k)+ ++ " has less items than "+ ++ "the other columns at index "+ ++ show i
+ src/DataFrame/Lazy/Internal/DataFrame.hs view
@@ -0,0 +1,92 @@+{-# LANGUAGE GADTs #-}+{-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE Strict #-}+{-# LANGUAGE BangPatterns #-}+{-# LANGUAGE OverloadedStrings #-}+{-# LANGUAGE NumericUnderscores #-}+module DataFrame.Lazy.Internal.DataFrame where++import Control.Monad (forM, foldM)+import Data.IORef+import Data.Kind+import qualified Data.List as L+import qualified Data.Map as M+import qualified Data.Text as T+import qualified Data.Vector as V+import qualified DataFrame.Internal.DataFrame as D+import qualified DataFrame.Internal.Column as C+import qualified DataFrame.Internal.Expression as E+import qualified DataFrame.Operations.Core as D+import DataFrame.Operations.Merge+import qualified DataFrame.Operations.Subset as D+import qualified DataFrame.Operations.Transformations as D+import qualified DataFrame.Lazy.IO.CSV as D+import System.FilePath++data LazyOperation where+ Derive :: C.Columnable a => T.Text -> E.Expr a -> LazyOperation+ Select :: [T.Text] -> LazyOperation+ Filter :: E.Expr Bool -> LazyOperation++instance Show LazyOperation where+ show :: LazyOperation -> String+ show (Derive name expr) = T.unpack name ++ " := " ++ show expr+ show (Select columns) = "select(" ++ show columns ++ ")"+ show (Filter expr) = "filter(" ++ show expr ++ ")"++data InputType = ICSV deriving Show++data LazyDataFrame = LazyDataFrame+ { inputPath :: FilePath+ , inputType :: InputType+ , operations :: [LazyOperation]+ , batchSize :: Int+ } deriving Show++eval :: LazyOperation -> D.DataFrame -> D.DataFrame+eval (Derive name expr) = D.derive name expr+eval (Select columns) = D.select columns+eval (Filter expr) = D.filterWhere expr++runDataFrame :: forall a . (C.Columnable a) => LazyDataFrame -> IO D.DataFrame+runDataFrame df = do+ let path = inputPath df+ totalRows <- D.countRows ',' path+ let batches = batchRanges totalRows (batchSize df)+ (df', _) <- foldM (\(!accDf, (!pos, !unused, !r)) (!start, !end) -> do+ mapM_ putStr ["Scanning: ", show start, " to ", show end, " rows out of ", show totalRows, "\n"] ++ (!sdf, (!pos', !unconsumed, !rowsRead)) <- D.readSeparated ',' (+ D.defaultOptions { D.rowRange = Just (start, batchSize df)+ , D.totalRows = Just totalRows+ , D.seekPos = pos+ , D.rowsRead = r+ , D.leftOver = unused}) path+ let !rdf = L.foldl' (flip eval) sdf (operations df)+ return (accDf <> rdf, (Just pos', unconsumed, rowsRead + r)) ) (D.empty, (Nothing, "", 0)) batches+ return df'++batchRanges :: Int -> Int -> [(Int, Int)]+batchRanges n inc = go n [0,inc..n]+ where+ go _ [] = []+ go n [x] = [(x, n)]+ go n (f:s:rest) =(f, s) : go n (s:rest)++scanCsv :: T.Text -> LazyDataFrame+scanCsv path = LazyDataFrame (T.unpack path) ICSV [] 512_000++addOperation :: LazyOperation -> LazyDataFrame -> LazyDataFrame+addOperation op df = df { operations = operations df ++ [op] }++derive :: C.Columnable a => T.Text -> E.Expr a -> LazyDataFrame -> LazyDataFrame+derive name expr = addOperation (Derive name expr)++select :: C.Columnable a => [T.Text] -> LazyDataFrame -> LazyDataFrame+select columns = addOperation (Select columns)++filter :: C.Columnable a => E.Expr Bool -> LazyDataFrame -> LazyDataFrame+filter cond = addOperation (Filter cond)
src/DataFrame/Operations/Aggregation.hs view
@@ -23,7 +23,7 @@ import Control.Exception (throw) import Control.Monad (foldM_) import Control.Monad.ST (runST)-import DataFrame.Internal.Column (Column(..), toColumn', getIndicesUnboxed, getIndices, Columnable)+import DataFrame.Internal.Column (Column(..), fromVector, getIndicesUnboxed, getIndices, Columnable) import DataFrame.Internal.DataFrame (DataFrame(..), empty, getColumn) import DataFrame.Internal.Parsing import DataFrame.Internal.Types@@ -181,10 +181,10 @@ DataFrame reduceBy f name df = case getColumn name df of Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> case testEquality (typeRep @a) (typeRep @a') of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map f column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map f column)) df Nothing -> error "Type error" Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> case testEquality (typeRep @a) (typeRep @a') of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map f column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map f column)) df Nothing -> error "Type error" _ -> error "Column is ungrouped" @@ -194,40 +194,40 @@ -> DataFrame reduceByAgg agg name df = case agg of Count -> case getColumn name df of- Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> insertColumn' name (Just $ toColumn' (VG.map VG.length column)) df- Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> insertColumn' name (Just $ toColumn' (VG.map VG.length column)) df+ Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> insertColumn' name (Just $ fromVector (VG.map VG.length column)) df+ Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> insertColumn' name (Just $ fromVector (VG.map VG.length column)) df _ -> error $ "Cannot count ungrouped Column: " ++ T.unpack name Mean -> case getColumn name df of Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> case testEquality (typeRep @a') (typeRep @Int) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map (SS.mean . VG.map fromIntegral) column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map (SS.mean . VG.map fromIntegral) column)) df Nothing -> case testEquality (typeRep @a') (typeRep @Double) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map SS.mean column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map SS.mean column)) df Nothing -> case testEquality (typeRep @a') (typeRep @Float) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map (SS.mean . VG.map realToFrac) column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map (SS.mean . VG.map realToFrac) column)) df Nothing -> error $ "Cannot get mean of non-numeric column: " ++ T.unpack name -- Not sure what to do with no numeric - return nothing??? Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> case testEquality (typeRep @a') (typeRep @Int) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map (SS.mean . VG.map fromIntegral) column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map (SS.mean . VG.map fromIntegral) column)) df Nothing -> case testEquality (typeRep @a') (typeRep @Double) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map SS.mean column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map SS.mean column)) df Nothing -> case testEquality (typeRep @a') (typeRep @Float) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map (SS.mean . VG.map realToFrac) column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map (SS.mean . VG.map realToFrac) column)) df Nothing -> error $ "Cannot get mean of non-numeric column: " ++ T.unpack name -- Not sure what to do with no numeric - return nothing??? Minimum -> case getColumn name df of- Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> insertColumn' name (Just $ toColumn' (VG.map VG.minimum column)) df- Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> insertColumn' name (Just $ toColumn' (VG.map VG.minimum column)) df+ Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> insertColumn' name (Just $ fromVector (VG.map VG.minimum column)) df+ Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> insertColumn' name (Just $ fromVector (VG.map VG.minimum column)) df Maximum -> case getColumn name df of- Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> insertColumn' name (Just $ toColumn' (VG.map VG.maximum column)) df- Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> insertColumn' name (Just $ toColumn' (VG.map VG.maximum column)) df+ Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> insertColumn' name (Just $ fromVector (VG.map VG.maximum column)) df+ Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> insertColumn' name (Just $ fromVector (VG.map VG.maximum column)) df Sum -> case getColumn name df of Just ((GroupedBoxedColumn (column :: V.Vector (V.Vector a')))) -> case testEquality (typeRep @a') (typeRep @Int) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map VG.sum column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map VG.sum column)) df Nothing -> case testEquality (typeRep @a') (typeRep @Double) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map VG.sum column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map VG.sum column)) df Nothing -> error $ "Cannot get sum of non-numeric column: " ++ T.unpack name -- Not sure what to do with no numeric - return nothing??? Just ((GroupedUnboxedColumn (column :: V.Vector (VU.Vector a')))) -> case testEquality (typeRep @a') (typeRep @Int) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map VG.sum column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map VG.sum column)) df Nothing -> case testEquality (typeRep @a') (typeRep @Double) of- Just Refl -> insertColumn' name (Just $ toColumn' (VG.map VG.sum column)) df+ Just Refl -> insertColumn' name (Just $ fromVector (VG.map VG.sum column)) df Nothing -> error $ "Cannot get sum of non-numeric column: " ++ T.unpack name -- Not sure what to do with no numeric - return nothing??? _ -> error "UNIMPLEMENTED"
src/DataFrame/Operations/Core.hs view
@@ -6,6 +6,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE BangPatterns #-}+{-# LANGUAGE Strict #-} module DataFrame.Operations.Core where import qualified Data.List as L@@ -19,7 +20,7 @@ import Control.Exception ( throw ) import DataFrame.Errors-import DataFrame.Internal.Column ( Column(..), toColumn', toColumn, columnLength, columnTypeString, expandColumn, Columnable)+import DataFrame.Internal.Column ( Column(..), fromVector, fromList, columnLength, columnTypeString, expandColumn, Columnable) import DataFrame.Internal.DataFrame (DataFrame(..), getColumn, null, empty) import DataFrame.Internal.Parsing (isNullish) import Data.Either@@ -50,7 +51,7 @@ -- | DataFrame to add column to DataFrame -> DataFrame-insertColumn name xs = insertColumn' name (Just (toColumn' xs))+insertColumn name xs = insertColumn' name (Just (fromVector xs)) {-# INLINE insertColumn #-} cloneColumn :: T.Text -> T.Text -> DataFrame -> DataFrame@@ -96,7 +97,7 @@ | r > 0 && l > r = let indexes = (map snd . L.sortBy (compare `on` snd). M.toList . columnIndices) d nonEmptyColumns = L.foldl' (\acc i -> acc ++ [maybe (error "Unexpected") (expandColumn diff) (columns d V.! i)]) [] indexes- in fromList (zip (columnNames d ++ [name]) (nonEmptyColumns ++ [column]))+ in fromNamedColumns (zip (columnNames d ++ [name]) (nonEmptyColumns ++ [column])) | otherwise = let (n:rest) = case freeIndices d of [] -> [VG.length (columns d)..(VG.length (columns d) * 2 - 1)]@@ -133,7 +134,7 @@ insertColumnWithDefault defaultValue name xs d = let (rows, _) = dataframeDimensions d values = xs V.++ V.replicate (rows - V.length xs) defaultValue- in insertColumn' name (Just $ toColumn' values) d+ in insertColumn' name (Just $ fromVector values) d -- TODO: Add existence check in rename. rename :: T.Text -> T.Text -> DataFrame -> DataFrame@@ -159,12 +160,12 @@ -- | O(n) Returns the number of non-null columns in the dataframe and the type associated -- with each column. columnInfo :: DataFrame -> DataFrame-columnInfo df = empty & insertColumn' "Column Name" (Just $! toColumn (map nameOfColumn infos))- & insertColumn' "# Non-null Values" (Just $! toColumn (map nonNullValues infos))- & insertColumn' "# Null Values" (Just $! toColumn (map nullValues infos))- & insertColumn' "# Partially parsed" (Just $! toColumn (map partiallyParsedValues infos))- & insertColumn' "# Unique Values" (Just $! toColumn (map uniqueValues infos))- & insertColumn' "Type" (Just $! toColumn (map typeOfColumn infos))+columnInfo df = empty & insertColumn' "Column Name" (Just $! fromList (map nameOfColumn infos))+ & insertColumn' "# Non-null Values" (Just $! fromList (map nonNullValues infos))+ & insertColumn' "# Null Values" (Just $! fromList (map nullValues infos))+ & insertColumn' "# Partially parsed" (Just $! fromList (map partiallyParsedValues infos))+ & insertColumn' "# Unique Values" (Just $! fromList (map uniqueValues infos))+ & insertColumn' "Type" (Just $! fromList (map typeOfColumn infos)) where infos = L.sortBy (compare `on` nonNullValues) (V.ifoldl' go [] (columns df)) :: [ColumnInfo] indexMap = M.fromList (map (\(a, b) -> (b, a)) $ M.toList (columnIndices df))@@ -213,32 +214,47 @@ _ -> 0 partiallyParsed _ = 0 -fromList :: [(T.Text, Column)] -> DataFrame-fromList = L.foldl' (\df (!name, !column) -> insertColumn' name (Just $! column) df) empty+fromNamedColumns :: [(T.Text, Column)] -> DataFrame+fromNamedColumns = L.foldl' (\df (!name, !column) -> insertColumn' name (Just $! column) df) empty -fromColumnList :: [Column] -> DataFrame-fromColumnList = fromList . zip (map (T.pack . show) [0..])+fromUnamedColumns :: [Column] -> DataFrame+fromUnamedColumns = fromNamedColumns . zip (map (T.pack . show) [0..]) -- | O (k * n) Counts the occurences of each value in a given column. valueCounts :: forall a. (Columnable a) => T.Text -> DataFrame -> [(a, Int)] valueCounts columnName df = case getColumn columnName df of- Nothing -> throw $ ColumnNotFoundException columnName "sortBy" (map fst $ M.toList $ columnIndices df)+ Nothing -> throw $ ColumnNotFoundException columnName "valueCounts" (map fst $ M.toList $ columnIndices df) Just (BoxedColumn (column' :: V.Vector c)) -> let column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column' in case (typeRep @a) `testEquality` (typeRep @c) of- Nothing -> throw $ TypeMismatchException (typeRep @a) (typeRep @c) columnName "valueCounts"+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext+ { userType = Right $ typeRep @a+ , expectedType = Right $ typeRep @c+ , errorColumnName = Just (T.unpack columnName)+ , callingFunctionName = Just "valueCounts"+ }) Just Refl -> M.toAscList column Just (OptionalColumn (column' :: V.Vector c)) -> let column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty column' in case (typeRep @a) `testEquality` (typeRep @c) of- Nothing -> throw $ TypeMismatchException (typeRep @a) (typeRep @c) columnName "valueCounts"+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext+ { userType = Right $ typeRep @a+ , expectedType = Right $ typeRep @c+ , errorColumnName = Just (T.unpack columnName)+ , callingFunctionName = Just "valueCounts"+ }) Just Refl -> M.toAscList column Just (UnboxedColumn (column' :: VU.Vector c)) -> let column = V.foldl' (\m v -> MS.insertWith (+) v (1 :: Int) m) M.empty (V.convert column') in case (typeRep @a) `testEquality` (typeRep @c) of- Nothing -> throw $ TypeMismatchException (typeRep @a) (typeRep @c) columnName "valueCounts"+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext+ { userType = Right $ typeRep @a+ , expectedType = Right $ typeRep @c+ , errorColumnName = Just (T.unpack columnName)+ , callingFunctionName = Just "valueCounts"+ }) Just Refl -> M.toAscList column fold :: (a -> DataFrame -> DataFrame) -> [a] -> DataFrame -> DataFrame
+ src/DataFrame/Operations/Merge.hs view
@@ -0,0 +1,39 @@+{-# LANGUAGE InstanceSigs #-}+{-# LANGUAGE Strict #-}+module DataFrame.Operations.Merge where++import qualified Data.List as L+import qualified Data.Text as T+import qualified Data.Vector as V+import qualified DataFrame.Internal.Column as D+import qualified DataFrame.Internal.DataFrame as D+import qualified DataFrame.Operations.Core as D++instance Semigroup D.DataFrame where+ (<>) :: D.DataFrame -> D.DataFrame -> D.DataFrame+ (<>) a b = let+ columnsInBOnly = filter (\c -> c `notElem` D.columnNames b) (D.columnNames b)+ columnsInA = D.columnNames a+ addColumns a' b' df name+ | fst (D.dimensions a') == 0 && fst (D.dimensions b') == 0 = df+ | fst (D.dimensions a') == 0 = D.insertColumn' name (D.getColumn name b') df+ | fst (D.dimensions b') == 0 = D.insertColumn' name (D.getColumn name a') df+ | otherwise = let+ numColumnsA = (fst $ D.dimensions a')+ numColumnsB = (fst $ D.dimensions b')+ numColumns = max numColumnsA numColumnsB+ optA = D.getColumn name a'+ optB = D.getColumn name b'+ in case optB of+ Nothing -> case optA of+ Nothing -> D.insertColumn' name (Just (D.fromList ([] :: [T.Text]))) df+ Just a'' -> D.insertColumn' name (Just (D.expandColumn numColumnsB a'')) df+ Just b'' -> case optA of+ Nothing -> D.insertColumn' name (Just (D.leftExpandColumn numColumnsA b'')) df+ Just a'' -> D.insertColumn' name (D.concatColumns a'' b'') df+ in L.foldl' (addColumns a b) D.empty (D.columnNames a `L.union` D.columnNames b)++instance Monoid D.DataFrame where+ mempty = D.empty++
src/DataFrame/Operations/Statistics.hs view
@@ -9,6 +9,7 @@ module DataFrame.Operations.Statistics where import qualified Data.List as L+import qualified Data.Map as M import qualified Data.Text as T import qualified Data.Vector.Generic as VG import qualified Data.Vector as V@@ -32,6 +33,7 @@ frequencies :: T.Text -> DataFrame -> DataFrame frequencies name df = case getColumn name df of+ Nothing -> throw $ ColumnNotFoundException name "frequencies" (map fst $ M.toList $ columnIndices df) Just ((BoxedColumn (column :: V.Vector a))) -> let counts = valueCounts @a name df total = P.sum $ map snd counts@@ -88,7 +90,7 @@ correlation first second df = do f <- _getColumnAsDouble first df s <- _getColumnAsDouble second df- return $ SS.correlation (VG.zip f s)+ return $ SS.correlation2 f s _getColumnAsDouble :: T.Text -> DataFrame -> Maybe (VU.Vector Double) _getColumnAsDouble name df = case getColumn name df of@@ -101,24 +103,24 @@ sum :: T.Text -> DataFrame -> Maybe Double sum name df = case getColumn name df of+ Nothing -> throw $ ColumnNotFoundException name "sum" (map fst $ M.toList $ columnIndices df) Just ((UnboxedColumn (column :: VU.Vector a'))) -> case testEquality (typeRep @a') (typeRep @Int) of Just Refl -> Just $ VG.sum (VU.map fromIntegral column) Nothing -> case testEquality (typeRep @a') (typeRep @Double) of Just Refl -> Just $ VG.sum column Nothing -> Nothing- Nothing -> Nothing applyStatistic :: (VU.Vector Double -> Double) -> T.Text -> DataFrame -> Maybe Double applyStatistic f name df = do column <- getColumn name df if columnTypeString column == "Double"- then safeReduceColumn f column+ then reduceColumn f column else do- matching <- asum [transform (fromIntegral :: Int -> Double) column,- transform (fromIntegral :: Integer -> Double) column,- transform (realToFrac :: Float -> Double) column,+ matching <- asum [mapColumn (fromIntegral :: Int -> Double) column,+ mapColumn (fromIntegral :: Integer -> Double) column,+ mapColumn (realToFrac :: Float -> Double) column, Just column ]- safeReduceColumn f matching+ reduceColumn f matching applyStatistics :: (VU.Vector Double -> VU.Vector Double) -> T.Text -> DataFrame -> Maybe (VU.Vector Double) applyStatistics f name df = case getColumn name df of@@ -132,7 +134,7 @@ _ -> Nothing summarize :: DataFrame -> DataFrame-summarize df = fold columnStats (columnNames df) (fromList [("Statistic", toColumn ["Mean" :: T.Text, "Minimum", "25%" ,"Median", "75%", "Max", "StdDev", "IQR", "Skewness"])])+summarize df = fold columnStats (columnNames df) (fromNamedColumns [("Statistic", fromList ["Mean" :: T.Text, "Minimum", "25%" ,"Median", "75%", "Max", "StdDev", "IQR", "Skewness"])]) where columnStats name d = if all isJust (stats name) then insertUnboxedColumn name (VU.fromList (map (roundTo 2 . fromMaybe 0) $ stats name)) d else d stats name = let quantiles = applyStatistics (SS.quantilesVec SS.medianUnbiased (VU.fromList [0,1,2,3,4]) 4) name df
src/DataFrame/Operations/Subset.hs view
@@ -17,7 +17,7 @@ import qualified Prelude import Control.Exception (throw)-import DataFrame.Errors (DataFrameException(..))+import DataFrame.Errors (DataFrameException(..), TypeErrorContext(..)) import DataFrame.Internal.Column import DataFrame.Internal.DataFrame (DataFrame(..), getColumn, empty) import DataFrame.Internal.Expression@@ -80,7 +80,11 @@ filter filterColumnName condition df = case getColumn filterColumnName df of Nothing -> throw $ ColumnNotFoundException filterColumnName "filter" (map fst $ M.toList $ columnIndices df) Just column -> case ifoldlColumn (\s i v -> if condition v then S.insert i s else s) S.empty column of- Nothing -> throw $ TypeMismatchException' (typeRep @a) (columnTypeString column) filterColumnName "filter"+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext+ { userType = Right $ typeRep @a+ , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) + , errorColumnName = Just (T.unpack filterColumnName)+ , callingFunctionName = Just "filter"}) Just indexes -> let c' = snd $ dataframeDimensions df pick idxs col = atIndices idxs <$> col@@ -99,10 +103,10 @@ filterWhere :: Expr Bool -> DataFrame -> DataFrame filterWhere expr df = let (TColumn col) = interpret @Bool df expr- (Just indexes) = ifoldlColumn (\s i satisfied -> if satisfied then S.insert i s else s) S.empty col+ (Just indexes) = VU.convert . V.map (fromMaybe 0) . V.filter isJust . toVector @(Maybe Int) <$> imapColumn (\i satisfied -> if satisfied then Just i else Nothing) col c' = snd $ dataframeDimensions df- pick idxs col = atIndices idxs <$> col- in df {columns = V.map (pick indexes) (columns df), dataframeDimensions = (S.size indexes, c')}+ pick idxs col = atIndicesStable idxs <$> col+ in df {columns = V.map (pick indexes) (columns df), dataframeDimensions = (VU.length indexes, c')} -- | O(k) removes all rows with `Nothing` in a given column from the dataframe.
src/DataFrame/Operations/Transformations.hs view
@@ -3,6 +3,8 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE FlexibleContexts #-}+{-# LANGUAGE Strict #-}+{-# LANGUAGE StrictData #-} module DataFrame.Operations.Transformations where import qualified Data.List as L@@ -13,14 +15,14 @@ import qualified Data.Vector.Unboxed as VU import Control.Exception (throw)-import DataFrame.Errors (DataFrameException(..))-import DataFrame.Internal.Column (Column(..), columnTypeString, itransform, ifoldrColumn, TypedColumn (TColumn), Columnable, transform, unwrapTypedColumn)+import DataFrame.Errors (DataFrameException(..), TypeErrorContext(..))+import DataFrame.Internal.Column (Column(..), columnTypeString, imapColumn, ifoldrColumn, TypedColumn (TColumn), Columnable, mapColumn, unwrapTypedColumn) import DataFrame.Internal.DataFrame (DataFrame(..), getColumn) import DataFrame.Internal.Expression import DataFrame.Internal.Row (mkRowFromArgs, RowValue, toRowValue) import DataFrame.Operations.Core import Data.Maybe-import Type.Reflection (typeRep, typeOf)+import Type.Reflection (typeRep, typeOf, TypeRep) -- | O(k) Apply a function to a given column in a dataframe. apply ::@@ -33,12 +35,32 @@ -- | DataFrame to apply operation to DataFrame -> DataFrame-apply f columnName d = case getColumn columnName d of- Nothing -> throw $ ColumnNotFoundException columnName "apply" (map fst $ M.toList $ columnIndices d)- Just column -> case transform f column of- Nothing -> throw $ TypeMismatchException' (typeRep @b) (columnTypeString column) columnName "apply"- column' -> insertColumn' columnName column' d+apply f columnName d = case safeApply f columnName d of+ Left exception -> throw exception+ Right df -> df +-- | O(k) Safe version of the apply function. Returns (instead of throwing) the error.+safeApply ::+ forall b c.+ (Columnable b, Columnable c) =>+ -- | function to apply+ (b -> c) ->+ -- | Column name+ T.Text ->+ -- | DataFrame to apply operation to+ DataFrame ->+ Either DataFrameException DataFrame+safeApply f columnName d = case getColumn columnName d of+ Nothing -> Left $ ColumnNotFoundException columnName "apply" (map fst $ M.toList $ columnIndices d)+ Just column -> case mapColumn f column of+ Nothing -> Left $ TypeMismatchException (MkTypeErrorContext+ { userType = Right $ typeRep @b+ , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) + , errorColumnName = Just (T.unpack columnName)+ , callingFunctionName = Just "apply"+ })+ column' -> Right $ insertColumn' columnName column' d+ -- | O(k) Apply a function to a combination of columns in a dataframe and -- add the result into `alias` column. derive :: forall a . Columnable a => T.Text -> Expr a -> DataFrame -> DataFrame@@ -96,7 +118,12 @@ applyWhere condition filterColumnName f columnName df = case getColumn filterColumnName df of Nothing -> throw $ ColumnNotFoundException filterColumnName "applyWhere" (map fst $ M.toList $ columnIndices df) Just column -> case ifoldrColumn (\i val acc -> if condition val then V.cons i acc else acc) V.empty column of- Nothing -> throw $ TypeMismatchException' (typeRep @a) (columnTypeString column) filterColumnName "applyWhere"+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext+ { userType = Right $ typeRep @a+ , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) + , errorColumnName = Just (T.unpack columnName)+ , callingFunctionName = Just "applyWhere"+ }) Just indexes -> if V.null indexes then df else L.foldl' (\d i -> applyAtIndex i f columnName d) df indexes@@ -116,8 +143,13 @@ DataFrame applyAtIndex i f columnName df = case getColumn columnName df of Nothing -> throw $ ColumnNotFoundException columnName "applyAtIndex" (map fst $ M.toList $ columnIndices df)- Just column -> case itransform (\index value -> if index == i then f value else value) column of- Nothing -> throw $ TypeMismatchException' (typeRep @a) (columnTypeString column) columnName "applyAtIndex"+ Just column -> case imapColumn (\index value -> if index == i then f value else value) column of+ Nothing -> throw $ TypeMismatchException (MkTypeErrorContext+ { userType = Right $ typeRep @a+ , expectedType = Left (columnTypeString column) :: Either String (TypeRep ()) + , errorColumnName = Just (T.unpack columnName)+ , callingFunctionName = Just "applyAtIndex"+ }) column' -> insertColumn' columnName column' df impute ::@@ -129,5 +161,8 @@ DataFrame impute columnName value df = case getColumn columnName df of Nothing -> throw $ ColumnNotFoundException columnName "impute" (map fst $ M.toList $ columnIndices df)- Just (OptionalColumn _) -> apply (fromMaybe value) columnName df+ Just (OptionalColumn _) -> case safeApply (fromMaybe value) columnName df of+ Left (TypeMismatchException context) -> throw $ TypeMismatchException (context { callingFunctionName = Just "impute" })+ Left exception -> throw exception+ Right res -> res _ -> error "Cannot impute to a non-Empty column"
tests/Main.hs view
@@ -25,8 +25,8 @@ import qualified Operations.Take testData :: D.DataFrame-testData = D.fromList [ ("test1", DI.toColumn ([1..26] :: [Int]))- , ("test2", DI.toColumn ['a'..'z'])+testData = D.fromNamedColumns [ ("test1", DI.fromList ([1..26] :: [Int]))+ , ("test2", DI.fromList ['a'..'z']) ] -- Dimensions@@ -45,19 +45,19 @@ parseDate :: Test parseDate = let expected = Just $ DI.BoxedColumn (V.fromList [fromGregorian 2020 02 14, fromGregorian 2021 02 14, fromGregorian 2022 02 14])- actual = D.parseDefault True $ Just $ DI.toColumn' (V.fromList ["2020-02-14" :: T.Text, "2021-02-14", "2022-02-14"])+ actual = D.parseDefault True $ Just $ DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "2021-02-14", "2022-02-14"]) in TestCase (assertEqual "Correctly parses gregorian date" expected actual) incompleteDataParseEither :: Test incompleteDataParseEither = let expected = Just $ DI.BoxedColumn (V.fromList [Right $ fromGregorian 2020 02 14, Left ("2021-02-" :: T.Text), Right $ fromGregorian 2022 02 14])- actual = D.parseDefault True $ Just $ DI.toColumn' (V.fromList ["2020-02-14" :: T.Text, "2021-02-", "2022-02-14"])+ actual = D.parseDefault True $ Just $ DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "2021-02-", "2022-02-14"]) in TestCase (assertEqual "Parses Either for gregorian date" expected actual) incompleteDataParseMaybe :: Test incompleteDataParseMaybe = let expected = Just $ DI.BoxedColumn (V.fromList [Just $ fromGregorian 2020 02 14, Nothing, Just $ fromGregorian 2022 02 14])- actual = D.parseDefault True $ Just $ DI.toColumn' (V.fromList ["2020-02-14" :: T.Text, "", "2022-02-14"])+ actual = D.parseDefault True $ Just $ DI.fromVector (V.fromList ["2020-02-14" :: T.Text, "", "2022-02-14"]) in TestCase (assertEqual "Parses Maybe for gregorian date with null/empty" expected actual) parseTests :: [Test]
tests/Operations/Apply.hs view
@@ -15,18 +15,18 @@ import Type.Reflection (typeRep) values :: [(T.Text, DI.Column)]-values = [ ("test1", DI.toColumn ([1..26] :: [Int]))- , ("test2", DI.toColumn (map show ['a'..'z']))- , ("test3", DI.toColumn ([1..26] :: [Int]))- , ("test4", DI.toColumn ['a'..'z'])- , ("test5", DI.toColumn ([1..26] :: [Int]))- , ("test6", DI.toColumn ['a'..'z'])- , ("test7", DI.toColumn ([1..26] :: [Int]))- , ("test8", DI.toColumn ['a'..'z'])+values = [ ("test1", DI.fromList ([1..26] :: [Int]))+ , ("test2", DI.fromList (map show ['a'..'z']))+ , ("test3", DI.fromList ([1..26] :: [Int]))+ , ("test4", DI.fromList ['a'..'z'])+ , ("test5", DI.fromList ([1..26] :: [Int]))+ , ("test6", DI.fromList ['a'..'z'])+ , ("test7", DI.fromList ([1..26] :: [Int]))+ , ("test8", DI.fromList ['a'..'z']) ] testData :: D.DataFrame-testData = D.fromList values+testData = D.fromNamedColumns values applyBoxedToUnboxed :: Test applyBoxedToUnboxed = TestCase (assertEqual "Boxed apply unboxed when result is unboxed"@@ -40,7 +40,7 @@ applyWrongType :: Test applyWrongType = TestCase (assertExpectException "[Error Case]"- (DE.typeMismatchError (typeRep @Char) (typeRep @[Char]))+ (DE.typeMismatchError (show $ typeRep @Char) (show $ typeRep @[Char])) (print $ DI.getColumn "test2" $ D.apply @Char (const (1::Int)) "test2" testData)) applyUnknownColumn :: Test@@ -50,7 +50,7 @@ applyManyOnlyGivenFields :: Test applyManyOnlyGivenFields = TestCase (assertEqual "Applies function to many fields"- (D.fromList (map (, D.toColumn $ replicate 26 (1 :: Integer)) ["test4", "test6"] +++ (D.fromNamedColumns (map (, D.fromList $ replicate 26 (1 :: Integer)) ["test4", "test6"] ++ -- All other fields should have their original values. filter (\(name, col) -> name /= "test4" && name /= "test6") values)) (D.applyMany @Char (const (1::Integer))@@ -58,13 +58,13 @@ applyManyBoxedToBoxed :: Test applyManyBoxedToBoxed = TestCase (assertEqual "Applies function to many fields"- (D.fromList (map (, D.toColumn $ replicate 26 (1 :: Integer)) ["test4", "test6", "test8"]))+ (D.fromNamedColumns (map (, D.fromList $ replicate 26 (1 :: Integer)) ["test4", "test6", "test8"])) (D.select ["test4", "test6", "test8"] $ D.applyMany @Char (const (1::Integer)) ["test4", "test6", "test8"] testData)) applyManyBoxedToUnboxed :: Test applyManyBoxedToUnboxed = TestCase (assertEqual "Unboxes fields when necessary"- (D.fromList (map (, D.toColumn $ replicate 26 (1 :: Int)) ["test4", "test6", "test8"]))+ (D.fromNamedColumns (map (, D.fromList $ replicate 26 (1 :: Int)) ["test4", "test6", "test8"])) (D.select ["test4", "test6", "test8"] $ D.applyMany @Char (const (1::Int)) ["test4", "test6", "test8"] testData)) @@ -76,17 +76,17 @@ applyManyWrongType :: Test applyManyWrongType = TestCase (assertExpectException "[Error Case]"- (DE.typeMismatchError (typeRep @Char) (typeRep @[Char]))+ (DE.typeMismatchError (show $ typeRep @Char) (show $ typeRep @[Char])) (print $ DI.getColumn "test2" $ D.applyMany @Char (const (1::Int)) ["test2"] testData)) applyWhereWrongConditionType :: Test applyWhereWrongConditionType = TestCase (assertExpectException "[Error Case]"- (DE.typeMismatchError (typeRep @Integer) (typeRep @Int))+ (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) (print $ D.applyWhere (even @Integer) "test1" ((+1) :: Int -> Int) "test5" testData)) applyWhereWrongTargetType :: Test applyWhereWrongTargetType = TestCase (assertExpectException "[Error Case]"- (DE.typeMismatchError (typeRep @Float) (typeRep @Int))+ (DE.typeMismatchError (show $ typeRep @Float) (show $ typeRep @Int)) (print $ D.applyWhere (even @Int) "test1" ((+1) :: Float -> Float) "test5" testData)) applyWhereConditionColumnNotFound :: Test
tests/Operations/Derive.hs view
@@ -15,13 +15,13 @@ import Type.Reflection (typeRep) values :: [(T.Text, DI.Column)]-values = [ ("test1", DI.toColumn ([1..26] :: [Int]))- , ("test2", DI.toColumn (map show ['a'..'z']))- , ("test3", DI.toColumn ['a'..'z'])+values = [ ("test1", DI.fromList ([1..26] :: [Int]))+ , ("test2", DI.fromList (map show ['a'..'z']))+ , ("test3", DI.fromList ['a'..'z']) ] testData :: D.DataFrame-testData = D.fromList values+testData = D.fromNamedColumns values deriveWAI :: Test deriveWAI = TestCase (assertEqual "derive works with column expression"
tests/Operations/Filter.hs view
@@ -14,18 +14,18 @@ import Type.Reflection (typeRep) values :: [(T.Text, DI.Column)]-values = [ ("test1", DI.toColumn ([1..26] :: [Int]))- , ("test2", DI.toColumn (map show ['a'..'z']))- , ("test3", DI.toColumn ([1..26] :: [Int]))- , ("test4", DI.toColumn ['a'..'z'])- , ("test5", DI.toColumn ([1..26] :: [Int]))- , ("test6", DI.toColumn ['a'..'z'])- , ("test7", DI.toColumn ([1..26] :: [Int]))- , ("test8", DI.toColumn ['a'..'z'])+values = [ ("test1", DI.fromList ([1..26] :: [Int]))+ , ("test2", DI.fromList (map show ['a'..'z']))+ , ("test3", DI.fromList ([1..26] :: [Int]))+ , ("test4", DI.fromList ['a'..'z'])+ , ("test5", DI.fromList ([1..26] :: [Int]))+ , ("test6", DI.fromList ['a'..'z'])+ , ("test7", DI.fromList ([1..26] :: [Int]))+ , ("test8", DI.fromList ['a'..'z']) ] testData :: D.DataFrame-testData = D.fromList values+testData = D.fromNamedColumns values filterColumnDoesNotExist :: Test filterColumnDoesNotExist = TestCase (assertExpectException "[Error Case]"@@ -34,7 +34,7 @@ filterColumnWrongType :: Test filterColumnWrongType = TestCase (assertExpectException "[Error Case]"- (DE.typeMismatchError (typeRep @Integer) (typeRep @Int))+ (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) (print $ D.filter @Integer "test1" even testData)) filterByColumnDoesNotExist :: Test@@ -44,7 +44,7 @@ filterByColumnWrongType :: Test filterByColumnWrongType = TestCase (assertExpectException "[Error Case]"- (DE.typeMismatchError (typeRep @Integer) (typeRep @Int))+ (DE.typeMismatchError (show $ typeRep @Integer) (show $ typeRep @Int)) (print $ D.filterBy @Integer even "test1" testData)) filterColumnInexistentValues :: Test@@ -59,8 +59,8 @@ filterJustWAI :: Test filterJustWAI = TestCase (assertEqual "Filters out Nothing and unwraps Maybe"- (D.fromList [("test", D.toColumn $ replicate 5 (1 :: Int))])- (D.filterJust "test" (D.fromList [("test", D.toColumn $ take 10 $ cycle [Just (1 :: Int), Nothing])])))+ (D.fromNamedColumns [("test", D.fromList $ replicate 5 (1 :: Int))])+ (D.filterJust "test" (D.fromNamedColumns [("test", D.fromList$ take 10 $ cycle [Just (1 :: Int), Nothing])]))) tests :: [Test] tests = [ TestLabel "filterColumnDoesNotExist" filterColumnDoesNotExist
tests/Operations/GroupBy.hs view
@@ -12,50 +12,50 @@ import Test.HUnit values :: [(T.Text, DI.Column)]-values = [ ("test1", DI.toColumn (concatMap (replicate 10) [1 :: Int, 2, 3, 4]))- , ("test2", DI.toColumn (take 40 $ cycle [1 :: Int,2]))- , ("test3", DI.toColumn [(1 :: Int)..40])- , ("test4", DI.toColumn (reverse [(1 :: Int)..40]))+values = [ ("test1", DI.fromList (concatMap (replicate 10) [1 :: Int, 2, 3, 4]))+ , ("test2", DI.fromList (take 40 $ cycle [1 :: Int,2]))+ , ("test3", DI.fromList [(1 :: Int)..40])+ , ("test4", DI.fromList (reverse [(1 :: Int)..40])) ] testData :: D.DataFrame-testData = D.fromList values+testData = D.fromNamedColumns values groupBySingleRowWAI :: Test groupBySingleRowWAI = TestCase (assertEqual "Groups by single column"- (D.fromList [("test1", DI.toColumn [(1::Int)..4]),- -- This just makes rows with [1, 2] for every unique test1 row- ("test2", DI.GroupedUnboxedColumn (V.replicate 4 $ VU.fromList (take 10 $ cycle [1 :: Int, 2]))),- ("test3", DI.GroupedUnboxedColumn (V.generate 4 (\i -> VU.fromList [(i * 10 + 1)..((i + 1) * 10)]))),- ("test4", DI.GroupedUnboxedColumn (V.generate 4 (\i -> VU.fromList [(((3 - i) + 1) * 10),(((3 - i) + 1) * 10 - 1)..((3 - i) * 10 + 1)])))- ])+ (D.fromNamedColumns [("test1", DI.fromList [(1::Int)..4]),+ -- This just makes rows with [1, 2] for every unique test1 row+ ("test2", DI.GroupedUnboxedColumn (V.replicate 4 $ VU.fromList (take 10 $ cycle [1 :: Int, 2]))),+ ("test3", DI.GroupedUnboxedColumn (V.generate 4 (\i -> VU.fromList [(i * 10 + 1)..((i + 1) * 10)]))),+ ("test4", DI.GroupedUnboxedColumn (V.generate 4 (\i -> VU.fromList [(((3 - i) + 1) * 10),(((3 - i) + 1) * 10 - 1)..((3 - i) * 10 + 1)])))+ ]) (D.groupBy ["test1"] testData D.|> D.sortBy D.Ascending ["test1"])) groupByMultipleRowsWAI :: Test groupByMultipleRowsWAI = TestCase (assertEqual "Groups by single column"- (D.fromList [("test1", DI.toColumn $ concatMap (replicate 2) [(1::Int)..4]),- ("test2", DI.toColumn (take 8 $ cycle [1 :: Int, 2])),- ("test3", DI.GroupedUnboxedColumn (V.fromList [- VU.fromList [1 :: Int,3..9],- VU.fromList [2,4..10],- VU.fromList [11,13..19],- VU.fromList [12,14..20],- VU.fromList [21,23..29],- VU.fromList [22,24..30],- VU.fromList [31,33..39],- VU.fromList [32,34..40]- ])),- ("test4", DI.GroupedUnboxedColumn (V.fromList $ reverse [- VU.fromList [1 :: Int,3..9],- VU.fromList [2,4..10],- VU.fromList [11,13..19],- VU.fromList [12,14..20],- VU.fromList [21,23..29],- VU.fromList [22,24..30],- VU.fromList [31,33..39],- VU.fromList [32,34..40]- ]))- ])+ (D.fromNamedColumns [("test1", DI.fromList$ concatMap (replicate 2) [(1::Int)..4]),+ ("test2", DI.fromList (take 8 $ cycle [1 :: Int, 2])),+ ("test3", DI.GroupedUnboxedColumn (V.fromList [+ VU.fromList [1 :: Int,3..9],+ VU.fromList [2,4..10],+ VU.fromList [11,13..19],+ VU.fromList [12,14..20],+ VU.fromList [21,23..29],+ VU.fromList [22,24..30],+ VU.fromList [31,33..39],+ VU.fromList [32,34..40]+ ])),+ ("test4", DI.GroupedUnboxedColumn (V.fromList $ reverse [+ VU.fromList [1 :: Int,3..9],+ VU.fromList [2,4..10],+ VU.fromList [11,13..19],+ VU.fromList [12,14..20],+ VU.fromList [21,23..29],+ VU.fromList [22,24..30],+ VU.fromList [31,33..39],+ VU.fromList [32,34..40]+ ]))+ ]) (D.groupBy ["test1", "test2"] testData D.|> D.sortBy D.Ascending ["test1", "test2"])) groupByColumnDoesNotExist :: Test
tests/Operations/InsertColumn.hs view
@@ -12,14 +12,14 @@ import Test.HUnit testData :: D.DataFrame-testData = D.fromList [ ("test1", DI.toColumn ([1..26] :: [Int]))- , ("test2", DI.toColumn ['a'..'z'])- , ("test3", DI.toColumn ([1..26] :: [Int]))- , ("test4", DI.toColumn ['a'..'z'])- , ("test5", DI.toColumn ([1..26] :: [Int]))- , ("test6", DI.toColumn ['a'..'z'])- , ("test7", DI.toColumn ([1..26] :: [Int]))- , ("test8", DI.toColumn ['a'..'z'])+testData = D.fromNamedColumns [ ("test1", DI.fromList([1..26] :: [Int]))+ , ("test2", DI.fromList['a'..'z'])+ , ("test3", DI.fromList([1..26] :: [Int]))+ , ("test4", DI.fromList['a'..'z'])+ , ("test5", DI.fromList([1..26] :: [Int]))+ , ("test6", DI.fromList['a'..'z'])+ , ("test7", DI.fromList([1..26] :: [Int]))+ , ("test8", DI.fromList['a'..'z']) ] -- Adding a boxed vector to an empty dataframe creates a new column boxed containing the vector elements.@@ -30,8 +30,8 @@ addBoxedColumn' :: Test addBoxedColumn' = TestCase (assertEqual "Two columns should be equal"- (Just $ DI.toColumn ["Thuba" :: T.Text, "Zodwa", "Themba"])- (DI.getColumn "new" $ D.insertColumn' "new" (Just $ DI.toColumn ["Thuba" :: T.Text, "Zodwa", "Themba"]) D.empty))+ (Just $ DI.fromList["Thuba" :: T.Text, "Zodwa", "Themba"])+ (DI.getColumn "new" $ D.insertColumn' "new" (Just $ DI.fromList["Thuba" :: T.Text, "Zodwa", "Themba"]) D.empty)) -- Adding an boxed vector with an unboxable type (Int/Double) to an empty dataframe creates a new column boxed containing the vector elements. addUnboxedColumn :: Test@@ -41,8 +41,8 @@ addUnboxedColumn' :: Test addUnboxedColumn' = TestCase (assertEqual "Value should be boxed"- (Just $ DI.toColumn [1 :: Int, 2, 3])- (DI.getColumn "new" $ D.insertColumn' "new" (Just $ DI.toColumn [1 :: Int, 2, 3]) D.empty))+ (Just $ DI.fromList[1 :: Int, 2, 3])+ (DI.getColumn "new" $ D.insertColumn' "new" (Just $ DI.fromList[1 :: Int, 2, 3]) D.empty)) -- Adding a column with less values than the current DF dimensions adds column with optionals. addSmallerColumnBoxed :: Test@@ -76,16 +76,16 @@ addLargerColumnBoxed :: Test addLargerColumnBoxed = TestCase (assertEqual "Smaller lists should grow and contain optionals"- (D.fromList [("new", D.toColumn [Just "a" :: Maybe T.Text, Just "b", Just "c", Nothing, Nothing]),- ("newer", D.toColumn ["a" :: T.Text, "b", "c", "d", "e"])])+ (D.fromNamedColumns [("new", D.fromList [Just "a" :: Maybe T.Text, Just "b", Just "c", Nothing, Nothing]),+ ("newer", D.fromList ["a" :: T.Text, "b", "c", "d", "e"])]) (D.insertColumn "newer" (V.fromList ["a" :: T.Text, "b", "c", "d", "e"]) $ D.insertColumn "new" (V.fromList ["a" :: T.Text, "b", "c"]) D.empty)) addLargerColumnUnboxed :: Test addLargerColumnUnboxed = TestCase (assertEqual "Smaller lists should grow and contain optionals"- (D.fromList [("old", D.toColumn [Just 1 :: Maybe Int, Just 2, Nothing, Nothing, Nothing]),- ("new", D.toColumn [Just 1 :: Maybe Int, Just 2, Just 3, Nothing, Nothing]),- ("newer", D.toColumn [1 :: Int, 2, 3, 4, 5])])+ (D.fromNamedColumns [("old", D.fromList [Just 1 :: Maybe Int, Just 2, Nothing, Nothing, Nothing]),+ ("new", D.fromList [Just 1 :: Maybe Int, Just 2, Just 3, Nothing, Nothing]),+ ("newer", D.fromList [1 :: Int, 2, 3, 4, 5])]) (D.insertColumn "newer" (V.fromList [1 :: Int, 2, 3, 4, 5]) $ D.insertColumn "new" (V.fromList [1 :: Int, 2, 3]) $ D.insertColumn "old" (V.fromList [1 :: Int, 2]) D.empty))
tests/Operations/Sort.hs view
@@ -18,23 +18,23 @@ values :: [(T.Text, DI.Column)] values = let ns = shuffle' [(1::Int)..26] 26 $ mkStdGen 252- in [ ("test1", DI.toColumn ns)- , ("test2", DI.toColumn (map (chr . (+96)) ns))+ in [ ("test1", DI.fromList ns)+ , ("test2", DI.fromList (map (chr . (+96)) ns)) ] testData :: D.DataFrame-testData = D.fromList values+testData = D.fromNamedColumns values sortByAscendingWAI :: Test sortByAscendingWAI = TestCase (assertEqual "Sorting rows by ascending works as intended"- (D.fromList [("test1", DI.toColumn [(1::Int)..26]),- ("test2", DI.toColumn ['a'..'z'])])+ (D.fromNamedColumns [("test1", DI.fromList [(1::Int)..26]),+ ("test2", DI.fromList ['a'..'z'])]) (D.sortBy D.Ascending ["test1"] testData)) sortByDescendingWAI :: Test sortByDescendingWAI = TestCase (assertEqual "Sorting rows by descending works as intended"- (D.fromList [("test1", DI.toColumn $ reverse [(1::Int)..26]),- ("test2", DI.toColumn $ reverse ['a'..'z'])])+ (D.fromNamedColumns [("test1", DI.fromList $ reverse [(1::Int)..26]),+ ("test2", DI.fromList $ reverse ['a'..'z'])]) (D.sortBy D.Descending ["test1"] testData)) sortByColumnDoesNotExist :: Test
tests/Operations/Take.hs view
@@ -7,16 +7,16 @@ import Test.HUnit testData :: D.DataFrame-testData = D.fromList [ ("test1", DI.toColumn ([1..26] :: [Int]))- , ("test2", DI.toColumn ['a'..'z'])+testData = D.fromNamedColumns [ ("test1", DI.fromList ([1..26] :: [Int]))+ , ("test2", DI.fromList ['a'..'z']) ] takeWAI :: Test-takeWAI = TestCase (assertEqual "Gets first 10 numbers" (Just $ D.toColumn [(1 :: Int)..10]) (D.getColumn "test1" $ D.take 10 testData))+takeWAI = TestCase (assertEqual "Gets first 10 numbers" (Just $ D.fromList [(1 :: Int)..10]) (D.getColumn "test1" $ D.take 10 testData)) takeLastWAI :: Test-takeLastWAI = TestCase (assertEqual "Gets first 10 numbers" (Just $ D.toColumn [(17 :: Int)..26]) (D.getColumn "test1" $ D.takeLast 10 testData))+takeLastWAI = TestCase (assertEqual "Gets first 10 numbers" (Just $ D.fromList [(17 :: Int)..26]) (D.getColumn "test1" $ D.takeLast 10 testData)) lengthEqualsTakeParam :: Test lengthEqualsTakeParam = TestCase (assertEqual "should be (5, 2)" (5, 2) (D.dimensions $ D.take 5 testData))