diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,20 @@ +Copyright (c) 2026 Michael Chavinda + +Permission is hereby granted, free of charge, to any person obtaining +a copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, +EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/dataframe-csv.cabal b/dataframe-csv.cabal new file mode 100644 --- /dev/null +++ b/dataframe-csv.cabal @@ -0,0 +1,42 @@ +cabal-version: 2.4 +name: dataframe-csv +version: 1.0.0.0 + +synopsis: CSV reader and writer for the dataframe ecosystem. +description: + @DataFrame.IO.CSV@ — streaming CSV read/write built on top of + @cassava@. Plus @DataFrame.IO.Internal.MutableColumn@, the + mutable-column ingest helpers shared by the lazy CSV reader. + +bug-reports: https://github.com/mchav/dataframe/issues +license: MIT +license-file: LICENSE +author: Michael Chavinda +maintainer: mschavinda@gmail.com +copyright: (c) 2024-2025 Michael Chavinda +category: Data +tested-with: GHC ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.2 + +common warnings + ghc-options: + -Wincomplete-patterns + -Wincomplete-uni-patterns + -Wunused-imports + -Wunused-local-binds + +library + import: warnings + exposed-modules: + DataFrame.IO.CSV + DataFrame.IO.Internal.MutableColumn + build-depends: base >= 4 && < 5, + bytestring >= 0.11 && < 0.13, + cassava >= 0.1 && < 1, + containers >= 0.6.7 && < 0.9, + dataframe-core ^>= 1.0, + dataframe-operations ^>= 1.0, + dataframe-parsing ^>= 1.0, + text >= 2.0 && < 3, + vector ^>= 0.13 + hs-source-dirs: src + default-language: Haskell2010 diff --git a/src/DataFrame/IO/CSV.hs b/src/DataFrame/IO/CSV.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/CSV.hs @@ -0,0 +1,797 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE NumericUnderscores #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +module DataFrame.IO.CSV where + +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as C +import qualified Data.ByteString.Lazy as BL +import qualified Data.Map.Strict as M +import qualified Data.Proxy as P +import qualified Data.Text as T +import qualified Data.Text.Encoding as TE +import qualified Data.Text.IO as TIO +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as VM +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +import Data.Csv.Streaming (Records (..)) +import qualified Data.Csv.Streaming as CsvStream + +import Control.Exception (SomeException, catch) +import Control.Monad +import Control.Monad.ST (runST) +import Data.Char +import qualified Data.Csv as Csv +import Data.Either +import Data.Functor +import Data.IORef +import Data.Maybe +import Data.Type.Equality (TestEquality (testEquality)) +import Data.Word (Word8) +import DataFrame.Internal.Column +import DataFrame.Internal.DataFrame ( + DataFrame (..), + forceDataFrame, + toSeparated, + ) +import DataFrame.Internal.Parsing +import DataFrame.Internal.Schema +import DataFrame.Operations.Typing +import System.IO +import Type.Reflection +import Prelude hiding (concat, takeWhile) + +chunkSize :: Int +chunkSize = 16_384 + +data PagedVector a = PagedVector + { pvChunks :: !(IORef [V.Vector a]) + -- ^ Finished chunks (reverse order) + , pvActive :: !(IORef (VM.IOVector a)) + -- ^ Current mutable chunk + , pvCount :: !(IORef Int) + -- ^ Items written in current chunk + } + +data PagedUnboxedVector a = PagedUnboxedVector + { puvChunks :: !(IORef [VU.Vector a]) + , puvActive :: !(IORef (VUM.IOVector a)) + , puvCount :: !(IORef Int) + } + +data BuilderColumn + = BuilderInt !(PagedUnboxedVector Int) !(PagedUnboxedVector Word8) + | BuilderDouble !(PagedUnboxedVector Double) !(PagedUnboxedVector Word8) + | BuilderText !(PagedVector T.Text) !(PagedUnboxedVector Word8) + | BuilderBS !(PagedVector BS.ByteString) !(PagedUnboxedVector Word8) + +newPagedVector :: IO (PagedVector a) +newPagedVector = do + active <- VM.unsafeNew chunkSize + PagedVector <$> newIORef [] <*> newIORef active <*> newIORef 0 + +newPagedUnboxedVector :: (VUM.Unbox a) => IO (PagedUnboxedVector a) +newPagedUnboxedVector = do + active <- VUM.unsafeNew chunkSize + PagedUnboxedVector <$> newIORef [] <*> newIORef active <*> newIORef 0 + +appendPagedVector :: PagedVector a -> a -> IO () +appendPagedVector (PagedVector chunksRef activeRef countRef) !val = do + count <- readIORef countRef + active <- readIORef activeRef + + if count < chunkSize + then do + VM.unsafeWrite active count val + writeIORef countRef $! count + 1 + else do + frozen <- V.unsafeFreeze active + modifyIORef' chunksRef (frozen :) + + newActive <- VM.unsafeNew chunkSize + VM.unsafeWrite newActive 0 val + + writeIORef activeRef newActive + writeIORef countRef 1 +{-# INLINE appendPagedVector #-} + +appendPagedUnboxedVector :: (VUM.Unbox a) => PagedUnboxedVector a -> a -> IO () +appendPagedUnboxedVector (PagedUnboxedVector chunksRef activeRef countRef) !val = do + count <- readIORef countRef + active <- readIORef activeRef + + if count < chunkSize + then do + VUM.unsafeWrite active count val + writeIORef countRef $! count + 1 + else do + frozen <- VU.unsafeFreeze active + modifyIORef' chunksRef (frozen :) + + newActive <- VUM.unsafeNew chunkSize + VUM.unsafeWrite newActive 0 val + + writeIORef activeRef newActive + writeIORef countRef 1 +{-# INLINE appendPagedUnboxedVector #-} + +freezePagedVector :: PagedVector a -> IO (V.Vector a) +freezePagedVector (PagedVector chunksRef activeRef countRef) = do + count <- readIORef countRef + active <- readIORef activeRef + chunks <- readIORef chunksRef + + writeIORef chunksRef [] -- release chunk references + let frozenChunks = reverse chunks + totalLen = count + sum (map V.length frozenChunks) + + mv <- VM.unsafeNew totalLen + + let copyChunk !offset chunk = do + V.copy (VM.slice offset (V.length chunk) mv) chunk + pure (offset + V.length chunk) + + offset <- foldM copyChunk 0 frozenChunks + VM.copy (VM.slice offset count mv) (VM.slice 0 count active) + + V.unsafeFreeze mv + +freezePagedUnboxedVector :: + (VUM.Unbox a) => PagedUnboxedVector a -> IO (VU.Vector a) +freezePagedUnboxedVector (PagedUnboxedVector chunksRef activeRef countRef) = do + count <- readIORef countRef + active <- readIORef activeRef + chunks <- readIORef chunksRef + + writeIORef chunksRef [] -- release chunk references + let frozenChunks = reverse chunks + totalLen = count + sum (map VU.length frozenChunks) + + mv <- VUM.unsafeNew totalLen + + let copyChunk !offset chunk = do + VU.copy (VUM.slice offset (VU.length chunk) mv) chunk + pure (offset + VU.length chunk) + + offset <- foldM copyChunk 0 frozenChunks + VUM.copy (VUM.slice offset count mv) (VUM.slice 0 count active) + + VU.unsafeFreeze mv + +-- | STANDARD CONFIG TYPES +data HeaderSpec = NoHeader | UseFirstRow | ProvideNames [T.Text] + deriving (Eq, Show) + +data TypeSpec + = InferFromSample Int + | SpecifyTypes [(T.Text, SchemaType)] TypeSpec + | NoInference + +{- | How the fast reader should treat a row whose field count does not +match the header row. Only consulted by @dataframe-fastcsv@; the pure +Haskell reader has its own semantics. +-} +data RaggedRowPolicy + = {- | Fill missing cells with nulls; silently drop extras. + Matches pandas / polars lenient defaults. + -} + PadWithNull + | {- | Fill missing cells with nulls; silently drop extras. Alias + kept for ergonomic naming when the caller only cares about the + "don't raise, just forget the extras" half of 'PadWithNull'. + -} + Truncate + | {- | Raise a 'DataFrame.IO.CSV.Fast.CsvParseError' on any row whose + field count differs from the header. Matches polars' strict + schema-bound mode. + -} + RaiseOnRagged + deriving (Eq, Show) + +-- | How the fast reader should treat an unclosed quoted field at EOF. +data UnclosedQuotePolicy + = -- | Raise a 'DataFrame.IO.CSV.Fast.CsvParseError'. Default. + RaiseOnUnclosedQuote + | {- | Return whatever rows were parsed before the stray quote; the + remainder is silently dropped. + -} + BestEffort + deriving (Eq, Show) + +-- | CSV read parameters. +data ReadOptions = ReadOptions + { headerSpec :: HeaderSpec + -- ^ Where to get the headers from. (default: UseFirstRow) + , typeSpec :: TypeSpec + -- ^ Whether/how to infer types. (default: InferFromSample 100) + , safeRead :: SafeReadMode + {- ^ Default 'SafeReadMode' for columns without an entry in + 'safeReadOverrides'. (default: 'NoSafeRead') + -} + , safeReadOverrides :: [(T.Text, SafeReadMode)] + -- ^ Per-column 'SafeReadMode' overrides; takes precedence over 'safeRead'. + , dateFormat :: String + {- ^ Format of date fields as recognized by the Data.Time.Format module. + + __Examples:__ + + @ + > parseTimeM True defaultTimeLocale "%Y/%-m/%-d" "2010/3/04" :: Maybe Day + Just 2010-03-04 + > parseTimeM True defaultTimeLocale "%d/%-m/%-Y" "04/3/2010" :: Maybe Day + Just 2010-03-04 + @ + -} + , columnSeparator :: Char + -- ^ Character that separates column values. + , numColumns :: Maybe Int + -- ^ Number of columns to read. + , missingIndicators :: [T.Text] + -- ^ Values that should be read as `Nothing`. + , fastCsvOnRaggedRow :: RaggedRowPolicy + {- ^ @dataframe-fastcsv@: how to treat rows with a non-header field count. + (default: 'PadWithNull') + -} + , fastCsvOnUnclosedQuote :: UnclosedQuotePolicy + {- ^ @dataframe-fastcsv@: how to treat an unclosed quoted field at EOF. + (default: 'RaiseOnUnclosedQuote') + -} + , fastCsvTrimUnquoted :: Bool + {- ^ @dataframe-fastcsv@: if 'True', leading/trailing whitespace is + stripped from unquoted fields after decoding. RFC 4180 preserves + this whitespace, and that is the default ('False'). + -} + } + +shouldInferFromSample :: TypeSpec -> Bool +shouldInferFromSample (InferFromSample _) = True +shouldInferFromSample (SpecifyTypes _ fallback) = shouldInferFromSample fallback +shouldInferFromSample _ = False + +schemaTypeMap :: TypeSpec -> M.Map T.Text SchemaType +schemaTypeMap (SpecifyTypes xs _) = M.fromList xs +schemaTypeMap _ = M.empty + +typeInferenceSampleSize :: TypeSpec -> Int +typeInferenceSampleSize (InferFromSample n) = n +typeInferenceSampleSize (SpecifyTypes _ fallback) = typeInferenceSampleSize fallback +typeInferenceSampleSize _ = 0 + +defaultReadOptions :: ReadOptions +defaultReadOptions = + ReadOptions + { headerSpec = UseFirstRow + , typeSpec = InferFromSample 100 + , safeRead = NoSafeRead + , safeReadOverrides = [] + , dateFormat = "%Y-%m-%d" + , columnSeparator = ',' + , numColumns = Nothing + , missingIndicators = + ["Nothing", "NULL", "", " ", "nan", "null", "N/A", "NaN", "NAN", "NA"] + , fastCsvOnRaggedRow = PadWithNull + , fastCsvOnUnclosedQuote = RaiseOnUnclosedQuote + , fastCsvTrimUnquoted = False + } + +{- | Read CSV file from path and load it into a dataframe. + +==== __Example__ +@ +ghci> D.readCsv ".\/data\/taxi.csv" + +@ +-} +readCsv :: FilePath -> IO DataFrame +readCsv = readSeparated defaultReadOptions + +type CsvReader = Schema -> FilePath -> IO DataFrame + +{- | Schema-driven attoparsec CSV reader. Coerces each column to the +type declared in 'Schema'; columns absent from the schema fall back to +the default inference path. Defined in terms of 'readSeparated' with +the 'TypeSpec' filled in. + +@ +import qualified DataFrame as D +df <- D.readCsvWithSchema schema "input.csv" +@ +-} +readCsvWithSchema :: CsvReader +readCsvWithSchema schema = + readSeparated + defaultReadOptions + { typeSpec = + SpecifyTypes + (M.toList (elements schema)) + (typeSpec defaultReadOptions) + } + +{- | Read CSV file from path and load it into a dataframe. + +==== __Example__ +@ +ghci> D.readCsvWithOpts ".\/data\/taxi.csv" (D.defaultReadOptions { dateFormat = "%d/%-m/%-Y" }) + +@ +-} +readCsvWithOpts :: ReadOptions -> FilePath -> IO DataFrame +readCsvWithOpts = readSeparated + +{- | Read TSV (tab separated) file from path and load it into a dataframe. + +==== __Example__ +@ +ghci> D.readTsv ".\/data\/taxi.tsv" + +@ +-} +readTsv :: FilePath -> IO DataFrame +readTsv = readSeparated (defaultReadOptions{columnSeparator = '\t'}) + +{- | Read text file with specified delimiter into a dataframe. + +==== __Example__ +@ +ghci> D.readSeparated (D.defaultReadOptions { columnSeparator = ';' }) ".\/data\/taxi.txt" + +@ +-} +readSeparated :: ReadOptions -> FilePath -> IO DataFrame +readSeparated opts !path = do + let stripUtf8Bom bs = fromMaybe bs (BL.stripPrefix "\xEF\xBB\xBF" bs) + csvData <- stripUtf8Bom <$> BL.readFile path + fmap forceDataFrame (decodeSeparated opts csvData) + +decodeSeparated :: ReadOptions -> BL.ByteString -> IO DataFrame +decodeSeparated !opts csvData = do + let sep = columnSeparator opts + let decodeOpts = Csv.defaultDecodeOptions{Csv.decDelimiter = fromIntegral (ord sep)} + let stream = CsvStream.decodeWith decodeOpts Csv.NoHeader csvData + + let peekStream (Cons (Right row) rest) = return (row, rest) + peekStream (Cons (Left err) _) = error $ "Error parsing CSV header: " ++ err + peekStream (Nil Nothing _) = error "Empty CSV file" + peekStream (Nil (Just err) _) = error err + + (firstRowRaw, dataStream) <- peekStream stream + + let (columnNames, rowsToProcess) = case headerSpec opts of + NoHeader -> + ( map (T.pack . show) [0 .. V.length firstRowRaw - 1] + , Cons (Right firstRowRaw) dataStream + ) + UseFirstRow -> + ( map (T.strip . TE.decodeUtf8Lenient . BL.toStrict) (V.toList firstRowRaw) + , dataStream + ) + ProvideNames ns -> + ( ns ++ drop (length ns) (map (T.pack . show) [0 .. V.length firstRowRaw - 1]) + , Cons (Right firstRowRaw) dataStream + ) + + (sampleRow, _) <- peekStream rowsToProcess + builderCols <- initializeColumns columnNames (V.toList sampleRow) opts + let !builderColsV = V.fromList builderCols + let colNamesV = V.fromList columnNames + resolveMode = + effectiveSafeRead + (safeRead opts) + (safeReadOverrides opts) + -- If ANY column is EitherRead we keep every raw cell (including + -- "N/A" etc.) verbatim; otherwise the missing-indicator list applies. + anyEither = + any (\n -> resolveMode n == EitherRead) columnNames + missing = if anyEither then [] else missingIndicators opts + processStream missing rowsToProcess builderColsV (numColumns opts) + + frozenCols <- + V.zipWithM + (\name bc -> finalizeBuilderColumn (resolveMode name) opts bc) + colNamesV + builderColsV + let numRows = maybe 0 columnLength (frozenCols V.!? 0) + + let df = + DataFrame + frozenCols + (M.fromList (zip columnNames [0 ..])) + (numRows, V.length frozenCols) + M.empty -- TODO give typed column references + pure $ parseWithTypes resolveMode (schemaTypeMap (typeSpec opts)) df + +initializeColumns :: + [T.Text] -> [BL.ByteString] -> ReadOptions -> IO [BuilderColumn] +initializeColumns names _row opts = zipWithM initColumn names (map lookupType names) + where + typeMap = schemaTypeMap (typeSpec opts) + -- Return Nothing for columns that should be inferred from BS + shouldInfer = case typeSpec opts of + InferFromSample _ -> True + SpecifyTypes _ fallback -> shouldInferFromSample fallback + NoInference -> False + lookupType name = M.lookup name typeMap + resolveMode = + effectiveSafeRead (safeRead opts) (safeReadOverrides opts) + initColumn :: T.Text -> Maybe SchemaType -> IO BuilderColumn + initColumn name _ | resolveMode name == EitherRead = do + validityRef <- newPagedUnboxedVector + BuilderBS <$> newPagedVector <*> pure validityRef + initColumn _ Nothing | shouldInfer = do + validityRef <- newPagedUnboxedVector + BuilderBS <$> newPagedVector <*> pure validityRef + initColumn _ mtype = do + validityRef <- newPagedUnboxedVector + let t = fromMaybe (schemaType @T.Text) mtype + case t of + SType (_ :: P.Proxy a) -> case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> BuilderInt <$> newPagedUnboxedVector <*> pure validityRef + Nothing -> case testEquality (typeRep @a) (typeRep @Double) of + Just Refl -> BuilderDouble <$> newPagedUnboxedVector <*> pure validityRef + Nothing -> BuilderText <$> newPagedVector <*> pure validityRef + +processStream :: + [T.Text] -> + CsvStream.Records (V.Vector BL.ByteString) -> + V.Vector BuilderColumn -> + Maybe Int -> + IO () +processStream _ _ _ (Just 0) = return () +processStream missing (Cons (Right row) rest) cols n = + processRow missing row cols + >> processStream missing rest cols (fmap (flip (-) 1) n) +processStream _missing (Cons (Left err) _) _ _ = error ("CSV Parse Error: " ++ err) +processStream _missing (Nil _ _) _ _ = return () + +processRow :: + [T.Text] -> V.Vector BL.ByteString -> V.Vector BuilderColumn -> IO () +processRow missing !vals !cols = V.zipWithM_ processValue vals cols + where + processValue !bs !col = do + let !bs' = BL.toStrict bs + case col of + BuilderInt gv valid -> case readByteStringInt bs' of + Just !i -> appendPagedUnboxedVector gv i >> appendPagedUnboxedVector valid 1 + Nothing -> appendPagedUnboxedVector gv 0 >> appendPagedUnboxedVector valid 0 + BuilderDouble gv valid -> case readByteStringDouble bs' of + Just !d -> appendPagedUnboxedVector gv d >> appendPagedUnboxedVector valid 1 + Nothing -> appendPagedUnboxedVector gv 0.0 >> appendPagedUnboxedVector valid 0 + BuilderText gv valid -> do + let !val = T.strip (TE.decodeUtf8Lenient bs') + appendPagedVector gv val + appendPagedUnboxedVector valid (if val `elem` missing then 0 else 1) + BuilderBS gv valid -> do + let !bs'' = C.strip bs' + appendPagedVector gv bs'' + appendPagedUnboxedVector + valid + (if TE.decodeUtf8Lenient bs'' `elem` missing then 0 else 1) + +freezeBuilderColumn :: BuilderColumn -> IO Column +freezeBuilderColumn (BuilderInt gv validRef) = do + vec <- freezePagedUnboxedVector gv + valid <- freezePagedUnboxedVector validRef + if VU.all (== 1) valid + then return $! UnboxedColumn Nothing vec + else constructOptional vec valid +freezeBuilderColumn (BuilderDouble gv validRef) = do + vec <- freezePagedUnboxedVector gv + valid <- freezePagedUnboxedVector validRef + if VU.all (== 1) valid + then return $! UnboxedColumn Nothing vec + else constructOptional vec valid +freezeBuilderColumn (BuilderText gv validRef) = do + vec <- freezePagedVector gv + valid <- freezePagedUnboxedVector validRef + if VU.all (== 1) valid + then return $! BoxedColumn Nothing vec + else constructOptionalBoxed vec valid +freezeBuilderColumn (BuilderBS _ _) = + error + "freezeBuilderColumn: BuilderBS must be finalized via finalizeBuilderColumn" + +finalizeBuilderColumn :: + SafeReadMode -> ReadOptions -> BuilderColumn -> IO Column +finalizeBuilderColumn mode opts bc = do + col <- case bc of + BuilderBS gv validRef -> do + vec <- freezePagedVector gv + valid <- freezePagedUnboxedVector validRef + return $! inferColumnFromBS mode opts vec valid + _ -> freezeBuilderColumn bc + return $! case mode of + NoSafeRead -> col + MaybeRead -> ensureOptional col + EitherRead -> col + +inferColumnFromBS :: + SafeReadMode -> + ReadOptions -> + V.Vector BS.ByteString -> + VU.Vector Word8 -> + Column +inferColumnFromBS mode opts vec valid = + let sampleN = let n = typeInferenceSampleSize (typeSpec opts) in if n == 0 then 100 else n + dfmt = dateFormat opts + -- The sample IS still Maybe-wrapped; it's bounded at 100 rows + -- by default, so the allocation is ignorable. The previous + -- full-column `asMaybeFull = V.generate ...` allocation is + -- gone — handlers walk (vec, valid) directly. + samples = V.generate (min sampleN (V.length vec)) $ \i -> + if valid VU.! i == 1 then Just (vec V.! i) else Nothing + assumption = makeParsingAssumptionBS dfmt samples + in case mode of + EitherRead -> handleBSEither dfmt assumption vec valid + _ -> case assumption of + IntAssumption -> handleBSInt dfmt vec valid + DoubleAssumption -> handleBSDouble vec valid + BoolAssumption -> handleBSBool vec valid + DateAssumption -> handleBSDate dfmt vec valid + TextAssumption -> handleBSText vec valid + NoAssumption -> handleBSNo dfmt vec valid + +{- | 'EitherRead' wrap for the ByteString inference path: produce an +@Either Text a@ column. Raw input bytes are preserved verbatim; rows that were +marked invalid by the builder (e.g. empty cells before EitherRead disabled +missing-indicator detection) become @Left \"\"@. +-} +handleBSEither :: + String -> + ParsingAssumption -> + V.Vector BS.ByteString -> + VU.Vector Word8 -> + Column +handleBSEither dfmt assumption vec valid = case assumption of + BoolAssumption -> wrap readByteStringBool + IntAssumption -> wrap readByteStringInt + DoubleAssumption -> wrap readByteStringDouble + DateAssumption -> wrap (readByteStringDate dfmt) + -- Text / No assumption: column is Either Text Text; empty cells become + -- Left "" to keep the "Left means missing/failure" convention. + TextAssumption -> fromVector (V.imap textEither vec) + NoAssumption -> fromVector (V.imap textEither vec) + where + wrap :: + forall a. (Columnable a) => (BS.ByteString -> Maybe a) -> Column + wrap p = fromVector (V.imap (toEither p) vec) + + toEither :: + forall a. + (BS.ByteString -> Maybe a) -> + Int -> + BS.ByteString -> + Either T.Text a + toEither p i bs + | valid VU.! i == 0 = Left (TE.decodeUtf8Lenient bs) + | otherwise = case p bs of + Just v -> Right v + Nothing -> Left (TE.decodeUtf8Lenient bs) + + textEither :: Int -> BS.ByteString -> Either T.Text T.Text + textEither i bs = + let t = TE.decodeUtf8Lenient bs + in if valid VU.! i == 0 || T.null t then Left t else Right t + +makeParsingAssumptionBS :: + String -> V.Vector (Maybe BS.ByteString) -> ParsingAssumption +makeParsingAssumptionBS dfmt asMaybe + | V.all (== Nothing) asMaybe = NoAssumption + | vecSameConstructor asMaybe asMaybeBool = BoolAssumption + | vecSameConstructor asMaybe asMaybeInt + && vecSameConstructor asMaybe asMaybeDouble = + IntAssumption + | vecSameConstructor asMaybe asMaybeDouble = DoubleAssumption + | vecSameConstructor asMaybe asMaybeDate = DateAssumption + | otherwise = TextAssumption + where + asMaybeBool = V.map (>>= readByteStringBool) asMaybe + asMaybeInt = V.map (>>= readByteStringInt) asMaybe + asMaybeDouble = V.map (>>= readByteStringDouble) asMaybe + asMaybeDate = V.map (>>= readByteStringDate dfmt) asMaybe + +-- All @handleBS*@ helpers now take the raw @V.Vector BS.ByteString@ +-- plus the builder's @VU.Vector Word8@ validity vector, fusing the +-- parse + validity check into a single pass via +-- 'parseUnboxedColumnWithValid'. The previous @V.Vector (Maybe +-- BS.ByteString)@ intermediate — allocated upstream in +-- 'inferColumnFromBS' — is gone, along with the paired Int/Double +-- parses and the two 'V.zipWith' Bool vectors from +-- 'vecSameConstructor'. + +handleBSBool :: + V.Vector BS.ByteString -> VU.Vector Word8 -> Column +handleBSBool vec valid = + case parseUnboxedColumnWithValid False readByteStringBool vec valid of + Just (mbm, out) -> UnboxedColumn mbm out + Nothing -> handleBSText vec valid + +handleBSInt :: + String -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column +handleBSInt _dfmt vec valid = + case parseUnboxedColumnWithValid 0 readByteStringInt vec valid of + Just (mbm, out) -> UnboxedColumn mbm out + Nothing -> case parseUnboxedColumnWithValid 0 readByteStringDouble vec valid of + Just (mbm, out) -> UnboxedColumn mbm out + Nothing -> handleBSText vec valid + +handleBSDouble :: + V.Vector BS.ByteString -> VU.Vector Word8 -> Column +handleBSDouble vec valid = + case parseUnboxedColumnWithValid 0 readByteStringDouble vec valid of + Just (mbm, out) -> UnboxedColumn mbm out + Nothing -> handleBSText vec valid + +-- Dates are boxed ('Day' isn't 'VU.Unbox'), so fuse into a V.Vector +-- (Maybe Day) in one pass. Bails on the first non-null cell that +-- fails to parse. +handleBSDate :: + String -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column +handleBSDate dfmt vec valid = + case parseBoxedMaybeBSColumn valid (readByteStringDate dfmt) vec of + Just (anyNull, out) + | anyNull -> fromVector out + | otherwise -> fromVector (V.mapMaybe id out) + Nothing -> handleBSText vec valid + +-- Fused Text handler: decode each cell's UTF-8 bytes once, mark nulls +-- directly from the validity vector. Replaces the two-pass +-- `V.map (fmap decodeUtf8Lenient) ... sequenceA ...` pattern. +handleBSText :: + V.Vector BS.ByteString -> VU.Vector Word8 -> Column +handleBSText vec valid + | VU.any (== 0) valid = + fromVector + ( V.imap + ( \i bs -> + if valid VU.! i == 0 + then Nothing + else Just (TE.decodeUtf8Lenient bs) + ) + vec + ) + | otherwise = fromVector (V.map TE.decodeUtf8Lenient vec) + +handleBSNo :: + String -> V.Vector BS.ByteString -> VU.Vector Word8 -> Column +handleBSNo dfmt vec valid + | VU.all (== 0) valid = + fromVector (V.map (const (Nothing :: Maybe T.Text)) vec) + | Just (mbm, out) <- + parseUnboxedColumnWithValid False readByteStringBool vec valid = + UnboxedColumn mbm out + | Just (mbm, out) <- parseUnboxedColumnWithValid 0 readByteStringInt vec valid = + UnboxedColumn mbm out + | Just (mbm, out) <- parseUnboxedColumnWithValid 0 readByteStringDouble vec valid = + UnboxedColumn mbm out + | otherwise = case parseBoxedMaybeBSColumn valid (readByteStringDate dfmt) vec of + Just (anyNull, out) + | anyNull -> fromVector out + | otherwise -> fromVector (V.mapMaybe id out) + Nothing -> handleBSText vec valid + +-- Boxed counterpart to 'parseUnboxedColumnWithValid' for types that +-- aren't 'VU.Unbox' (e.g. 'Day'). Same one-pass + early-bail shape. +parseBoxedMaybeBSColumn :: + VU.Vector Word8 -> + (BS.ByteString -> Maybe a) -> + V.Vector BS.ByteString -> + Maybe (Bool, V.Vector (Maybe a)) +parseBoxedMaybeBSColumn valid parser vec = runST $ do + let n = V.length vec + out <- VM.new n + let loop !i !anyNull + | i >= n = do + frozen <- V.unsafeFreeze out + return (Just (anyNull, frozen)) + | VU.unsafeIndex valid i == 0 = do + VM.unsafeWrite out i Nothing + loop (i + 1) True + | otherwise = case parser (V.unsafeIndex vec i) of + Just v -> do + VM.unsafeWrite out i (Just v) + loop (i + 1) anyNull + Nothing -> return Nothing + loop 0 False + +{- | One-pass fused parse into a typed unboxed column. Avoids the +@V.Vector (Maybe a)@ intermediate that the "parse then 'sequenceA'" +idiom requires, and avoids the upstream @V.Vector (Maybe src)@ +classification by reading nullability from a precomputed validity +vector (produced by the CSV builder alongside the raw cells). + +Returns @Just (mbm, vec)@ only when every non-null cell parses. The +first unparseable cell short-circuits to @Nothing@ so the caller can +fall back to the next assumption (Int → Double → Text). @mbm@ is +@Nothing@ when no nulls exist (the column is non-nullable) and +@Just bm@ otherwise. + +Null slots are filled with @nullValue@; downstream consumers only see +them through the bitmap, so the sentinel never escapes. + +Memory shape for a length-@n@ input: + + * 1 × 'VUM.STVector' of @a@ (final data, @sizeOf a × n@ bytes) + * 1 × 'VUM.STVector' of 'Word8' (per-element validity, @n@ bytes) + * 1 × 'Bitmap' (bit-packed, @⌈n\/8⌉@ bytes) — only when nulls exist. +-} +parseUnboxedColumnWithValid :: + forall src a. + (VU.Unbox a) => + a -> + (src -> Maybe a) -> + V.Vector src -> + VU.Vector Word8 -> + Maybe (Maybe Bitmap, VU.Vector a) +parseUnboxedColumnWithValid nullValue parser vec valid = runST $ do + let n = V.length vec + values <- VUM.unsafeNew n + vmask <- VUM.unsafeNew n + let go !i !anyNull + | i >= n = finalizeParseResult values vmask anyNull + | VU.unsafeIndex valid i == 0 = do + VUM.unsafeWrite vmask i 0 + VUM.unsafeWrite values i nullValue + go (i + 1) True + | otherwise = case parser (V.unsafeIndex vec i) of + Just v -> do + VUM.unsafeWrite vmask i 1 + VUM.unsafeWrite values i v + go (i + 1) anyNull + Nothing -> return Nothing + go 0 False +{-# INLINE parseUnboxedColumnWithValid #-} + +constructOptional :: + (VU.Unbox a, Columnable a) => VU.Vector a -> VU.Vector Word8 -> IO Column +constructOptional vec valid = do + let bm = buildBitmapFromValid valid + pure $ UnboxedColumn (Just bm) vec + +constructOptionalBoxed :: V.Vector T.Text -> VU.Vector Word8 -> IO Column +constructOptionalBoxed vec valid = do + let bm = buildBitmapFromValid valid + pure $ BoxedColumn (Just bm) vec + +writeCsv :: FilePath -> DataFrame -> IO () +writeCsv = writeSeparated ',' + +writeTsv :: FilePath -> DataFrame -> IO () +writeTsv = writeSeparated '\t' + +writeSeparated :: + -- | Separator + Char -> + -- | Path to write to + FilePath -> + DataFrame -> + IO () +writeSeparated c filepath df = TIO.writeFile filepath (toSeparated c df) + +-- | Parse a CSV string into a DataFrame using default options. +fromCsv :: String -> IO (Either String DataFrame) +fromCsv s = do + let bs = BL.fromStrict (TE.encodeUtf8 (T.pack s)) + (Right <$> decodeSeparated defaultReadOptions bs) + `catch` (\(e :: SomeException) -> pure (Left (show e))) + +-- | Parse a lazy 'ByteString' containing CSV data into a DataFrame using default options. +fromCsvBytes :: BL.ByteString -> IO DataFrame +fromCsvBytes = decodeSeparated defaultReadOptions + +stripQuotes :: T.Text -> T.Text +stripQuotes txt = + case T.uncons txt of + Just ('"', rest) -> + case T.unsnoc rest of + Just (middle, '"') -> middle + _ -> txt + _ -> txt diff --git a/src/DataFrame/IO/Internal/MutableColumn.hs b/src/DataFrame/IO/Internal/MutableColumn.hs new file mode 100644 --- /dev/null +++ b/src/DataFrame/IO/Internal/MutableColumn.hs @@ -0,0 +1,89 @@ +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{- | +Mutable-column ingest helpers used by CSV readers. Kept out of +@DataFrame.Internal.Column@ so that the core column type does not need to +depend on @DataFrame.Internal.Parsing@. +-} +module DataFrame.IO.Internal.MutableColumn ( + writeColumn, + freezeColumn', +) where + +import qualified Data.Text as T +import qualified Data.Vector as VB +import qualified Data.Vector.Mutable as VBM +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +import Data.Maybe (fromMaybe) +import Data.Type.Equality (TestEquality (..), type (:~:) (Refl)) +import Type.Reflection (typeRep) + +import DataFrame.Internal.Column ( + Column (..), + MutableColumn (..), + buildBitmapFromNulls, + ) +import DataFrame.Internal.Parsing (isNullish, readDouble, readInt) + +writeColumn :: Int -> T.Text -> MutableColumn -> IO (Either T.Text Bool) +writeColumn i value (MBoxedColumn (col :: VBM.IOVector a)) = + case testEquality (typeRep @a) (typeRep @T.Text) of + Just Refl -> + if isNullish value + then VBM.unsafeWrite col i "" >> return (Left $! value) + else VBM.unsafeWrite col i value >> return (Right True) + Nothing -> return (Left value) +writeColumn i value (MUnboxedColumn (col :: VUM.IOVector a)) = + case testEquality (typeRep @a) (typeRep @Int) of + Just Refl -> case readInt value of + Just v -> VUM.unsafeWrite col i v >> return (Right True) + Nothing -> VUM.unsafeWrite col i 0 >> return (Left value) + Nothing -> case testEquality (typeRep @a) (typeRep @Double) of + Nothing -> return (Left $! value) + Just Refl -> case readDouble value of + Just v -> VUM.unsafeWrite col i v >> return (Right True) + Nothing -> VUM.unsafeWrite col i 0 >> return (Left $! value) +{-# INLINE writeColumn #-} + +freezeColumn' :: [(Int, T.Text)] -> MutableColumn -> IO Column +freezeColumn' nulls (MBoxedColumn col) + | null nulls = BoxedColumn Nothing <$> VB.unsafeFreeze col + | all (isNullish . snd) nulls = do + frozen <- VB.unsafeFreeze col + let n = VB.length frozen + bm = buildBitmapFromNulls n (map fst nulls) + return $ BoxedColumn (Just bm) frozen + | otherwise = + BoxedColumn Nothing + . VB.imap + ( \i v -> + if i `elem` map fst nulls + then Left (fromMaybe (error "UNEXPECTED ERROR DURING FREEZE") (lookup i nulls)) + else Right v + ) + <$> VB.unsafeFreeze col +freezeColumn' nulls (MUnboxedColumn col) + | null nulls = UnboxedColumn Nothing <$> VU.unsafeFreeze col + | all (isNullish . snd) nulls = do + c <- VU.unsafeFreeze col + let n = VU.length c + bm = buildBitmapFromNulls n (map fst nulls) + return $ UnboxedColumn (Just bm) c + | otherwise = do + c <- VU.unsafeFreeze col + return $ + BoxedColumn Nothing $ + VB.generate + (VU.length c) + ( \i -> + if i `elem` map fst nulls + then Left (fromMaybe (error "UNEXPECTED ERROR DURING FREEZE") (lookup i nulls)) + else Right (c VU.! i) + ) +{-# INLINE freezeColumn' #-}