diff --git a/README.md b/README.md --- a/README.md +++ b/README.md @@ -67,11 +67,14 @@ ### Visual example ![Screencast of usage in GHCI](./static/example.gif) +## 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. diff --git a/benchmark/Main.hs b/benchmark/Main.hs --- a/benchmark/Main.hs +++ b/benchmark/Main.hs @@ -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 diff --git a/dataframe.cabal b/dataframe.cabal --- a/dataframe.cabal +++ b/dataframe.cabal @@ -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 diff --git a/src/DataFrame.hs b/src/DataFrame.hs --- a/src/DataFrame.hs +++ b/src/DataFrame.hs @@ -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 diff --git a/src/DataFrame/Errors.hs b/src/DataFrame/Errors.hs --- a/src/DataFrame/Errors.hs +++ b/src/DataFrame/Errors.hs @@ -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 "" ) +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 :: )") + ++ green ("(" ++ cp ++ " ... :: )") ++ "\n\tor add " ++ "{-# LANGUAGE TypeApplications #-} to the top of your " ++ "file then change the call to \n\t\t" - ++ brightGreen (cp ++ " @ arg1 arg2") + ++ brightGreen (cp ++ " @ ....") guessColumnName :: T.Text -> [T.Text] -> T.Text guessColumnName userInput columns = case map (\k -> (editDistance userInput k, k)) columns of diff --git a/src/DataFrame/IO/Parquet.hs b/src/DataFrame/IO/Parquet.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Parquet.hs @@ -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)) diff --git a/src/DataFrame/Internal/Column.hs b/src/DataFrame/Internal/Column.hs --- a/src/DataFrame/Internal/Column.hs +++ b/src/DataFrame/Internal/Column.hs @@ -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}) diff --git a/src/DataFrame/Internal/DataFrame.hs b/src/DataFrame/Internal/DataFrame.hs --- a/src/DataFrame/Internal/DataFrame.hs +++ b/src/DataFrame/Internal/DataFrame.hs @@ -4,7 +4,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE GADTs #-} -{-# LANGUAGE StrictData #-} +{-# LANGUAGE Strict #-} {-# LANGUAGE FlexibleContexts #-} module DataFrame.Internal.DataFrame where diff --git a/src/DataFrame/Internal/Expression.hs b/src/DataFrame/Internal/Expression.hs --- a/src/DataFrame/Internal/Expression.hs +++ b/src/DataFrame/Internal/Expression.hs @@ -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 diff --git a/src/DataFrame/Internal/Types.hs b/src/DataFrame/Internal/Types.hs --- a/src/DataFrame/Internal/Types.hs +++ b/src/DataFrame/Internal/Types.hs @@ -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) diff --git a/src/DataFrame/Lazy.hs b/src/DataFrame/Lazy.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy.hs @@ -0,0 +1,3 @@ +module DataFrame.Lazy (module DataFrame.Lazy.Internal.DataFrame) where + +import DataFrame.Lazy.Internal.DataFrame diff --git a/src/DataFrame/Lazy/IO/CSV.hs b/src/DataFrame/Lazy/IO/CSV.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/IO/CSV.hs @@ -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 diff --git a/src/DataFrame/Lazy/Internal/DataFrame.hs b/src/DataFrame/Lazy/Internal/DataFrame.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Lazy/Internal/DataFrame.hs @@ -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) diff --git a/src/DataFrame/Operations/Aggregation.hs b/src/DataFrame/Operations/Aggregation.hs --- a/src/DataFrame/Operations/Aggregation.hs +++ b/src/DataFrame/Operations/Aggregation.hs @@ -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" diff --git a/src/DataFrame/Operations/Core.hs b/src/DataFrame/Operations/Core.hs --- a/src/DataFrame/Operations/Core.hs +++ b/src/DataFrame/Operations/Core.hs @@ -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 diff --git a/src/DataFrame/Operations/Merge.hs b/src/DataFrame/Operations/Merge.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/Operations/Merge.hs @@ -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 + + diff --git a/src/DataFrame/Operations/Statistics.hs b/src/DataFrame/Operations/Statistics.hs --- a/src/DataFrame/Operations/Statistics.hs +++ b/src/DataFrame/Operations/Statistics.hs @@ -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 diff --git a/src/DataFrame/Operations/Subset.hs b/src/DataFrame/Operations/Subset.hs --- a/src/DataFrame/Operations/Subset.hs +++ b/src/DataFrame/Operations/Subset.hs @@ -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. diff --git a/src/DataFrame/Operations/Transformations.hs b/src/DataFrame/Operations/Transformations.hs --- a/src/DataFrame/Operations/Transformations.hs +++ b/src/DataFrame/Operations/Transformations.hs @@ -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" diff --git a/tests/Main.hs b/tests/Main.hs --- a/tests/Main.hs +++ b/tests/Main.hs @@ -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] diff --git a/tests/Operations/Apply.hs b/tests/Operations/Apply.hs --- a/tests/Operations/Apply.hs +++ b/tests/Operations/Apply.hs @@ -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 diff --git a/tests/Operations/Derive.hs b/tests/Operations/Derive.hs --- a/tests/Operations/Derive.hs +++ b/tests/Operations/Derive.hs @@ -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" diff --git a/tests/Operations/Filter.hs b/tests/Operations/Filter.hs --- a/tests/Operations/Filter.hs +++ b/tests/Operations/Filter.hs @@ -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 diff --git a/tests/Operations/GroupBy.hs b/tests/Operations/GroupBy.hs --- a/tests/Operations/GroupBy.hs +++ b/tests/Operations/GroupBy.hs @@ -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 diff --git a/tests/Operations/InsertColumn.hs b/tests/Operations/InsertColumn.hs --- a/tests/Operations/InsertColumn.hs +++ b/tests/Operations/InsertColumn.hs @@ -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)) diff --git a/tests/Operations/Sort.hs b/tests/Operations/Sort.hs --- a/tests/Operations/Sort.hs +++ b/tests/Operations/Sort.hs @@ -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 diff --git a/tests/Operations/Take.hs b/tests/Operations/Take.hs --- a/tests/Operations/Take.hs +++ b/tests/Operations/Take.hs @@ -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))