dataframe-lazy-2.4.2.1: test/StreamingCsv.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

module Main (main) where

import Control.Exception (SomeException, bracket, try)
import Control.Monad (when)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as C8
import Data.IORef (modifyIORef', newIORef, readIORef)
import Data.List (isInfixOf)
import qualified Data.Map.Strict as M
import qualified Data.Text as T
import qualified DataFrame.IO.CSV as Csv
import DataFrame.Internal.Column (Column (..))
import qualified DataFrame.Internal.DataFrame as D
import qualified DataFrame.Lazy as L
import DataFrame.Schema (Schema (..), schemaType)
import System.Directory (removeFile)
import qualified System.Exit as Exit
import System.IO.Temp (emptySystemTempFile)
import Test.HUnit

withRawCsv :: BS.ByteString -> (FilePath -> IO a) -> IO a
withRawCsv bytes =
    bracket
        ( do
            path <- emptySystemTempFile "lazy_streaming_.csv"
            BS.writeFile path bytes
            pure path
        )
        removeFile

-- | Regression: a path-only reader sees one complete 64 MiB-or-EOF window.
streamingPathReaderFillsWindow :: Test
streamingPathReaderFillsWindow =
    TestLabel "streaming_path_reader_fills_window" $ TestCase $ do
        let row = "1," <> BS.replicate 4094 0x78 <> "\n"
            rowCount = 4096
            input = "id,payload\n" <> BS.concat (replicate rowCount row)
            schema = Schema $ M.fromList [("id", schemaType @Int)]
        withRawCsv input $ \path -> do
            calls <- newIORef (0 :: Int)
            let countingReader opts windowPath = do
                    modifyIORef' calls (+ 1)
                    Csv.readSeparated opts windowPath
            df <-
                L.runDataFrame
                    (L.scanCsvStreamingWith countingReader schema (T.pack path))
            assertEqual "all rows decoded" (rowCount, 1) (D.dataframeDimensions df)
            readIORef calls >>= assertEqual "one parser window" 1

-- | A record newline exactly at the 64 MiB read boundary must not be lost.
streamingWindowBoundaryIsLossless :: Test
streamingWindowBoundaryIsLossless =
    TestLabel "streaming_window_boundary_is_lossless" $ TestCase $ do
        let windowBytes = 64 * 1024 * 1024
            firstRecord = "1," <> BS.replicate (windowBytes - 3) 0x78 <> "\n"
            input = BS.concat ["id,payload\n", firstRecord, "2,y\n"]
            schema = Schema $ M.fromList [("id", schemaType @Int)]
        withRawCsv input $ \path -> do
            observations <- newIORef []
            let observingReader _ bytes = do
                    let (header, withNewline) = BS.break (== 0x0A) bytes
                        body = BS.drop 1 withNewline
                        observation =
                            ( C8.unpack header
                            , BS.length body
                            , BS.unpack (BS.take 2 body)
                            , if BS.null body then Nothing else Just (BS.last body)
                            )
                    modifyIORef' observations (observation :)
                    pure D.empty
            _ <-
                L.runDataFrame
                    (L.scanCsvStreamingBytesWith observingReader schema (T.pack path))
            actual <- reverse <$> readIORef observations
            assertEqual
                "two complete records, each with the original header"
                [ ("id,payload", windowBytes - 1, [0x31, 0x2C], Just 0x78)
                , ("id,payload", 3, [0x32, 0x2C], Just 0x79)
                ]
                actual

-- | A quoted newline spanning a 64 MiB boundary stays in one window.
streamingQuotedNewlineSpansWindow :: Test
streamingQuotedNewlineSpansWindow =
    TestLabel "streaming_quoted_newline_spans_window" $ TestCase $ do
        let windowBytes = 64 * 1024 * 1024
            firstChunk = "1,\"" <> BS.replicate (windowBytes - 4) 0x78 <> "\n"
            input = BS.concat ["id,payload\n", firstChunk, "\"\n"]
            schema = Schema $ M.fromList [("id", schemaType @Int)]
        withRawCsv input $ \path -> do
            bodyLengths <- newIORef []
            let observingReader _ bytes = do
                    let body = BS.drop 1 (snd (BS.break (== 0x0A) bytes))
                    modifyIORef' bodyLengths (BS.length body :)
                    pure D.empty
            _ <-
                L.runDataFrame
                    (L.scanCsvStreamingBytesWith observingReader schema (T.pack path))
            actual <- reverse <$> readIORef bodyLengths
            assertEqual
                "quoted record stays in one parser window"
                [windowBytes + 1]
                actual

-- | The default streaming path strips a BOM and matches an eager read.
streamingDefaultMatchesEager :: Test
streamingDefaultMatchesEager =
    TestLabel "streaming_default_matches_eager" $ TestCase $ do
        let input = "\xEF\xBB\xBFid,name,value\n1,Ada,3.5\n2,Alan,4.5"
            schema =
                Schema $
                    M.fromList
                        [ ("id", schemaType @Int)
                        , ("name", schemaType @T.Text)
                        , ("value", schemaType @Double)
                        ]
        withRawCsv input $ \path -> do
            let opts = Csv.schemaReadOptions schema
            eager <- Csv.readSeparated opts path
            streamed <- L.runDataFrame (L.scanCsvStreaming schema (T.pack path))
            assertEqual "streaming bytes reader matches eager reader" eager streamed

-- | A parser failure on the producer thread must fail the consumer.
streamingReaderFailurePropagates :: Test
streamingReaderFailurePropagates =
    TestLabel "streaming_reader_failure_propagates" $ TestCase $ do
        let input = "id\n1\n"
            schema = Schema $ M.fromList [("id", schemaType @Int)]
            failingReader _ _ = ioError (userError "streaming-reader-marker")
        withRawCsv input $ \path -> do
            result <-
                try
                    ( L.runDataFrame
                        (L.scanCsvStreamingBytesWith failingReader schema (T.pack path))
                    )
            case result of
                Left (err :: SomeException) ->
                    assertBool
                        ("unexpected exception: " <> show err)
                        ("streaming-reader-marker" `isInfixOf` show err)
                Right (_ :: D.DataFrame) -> assertFailure "expected reader failure"

-- | Schema-declared text stays in the CSV reader's packed representation.
streamingKeepsTextPacked :: Test
streamingKeepsTextPacked =
    TestLabel "streaming_keeps_text_packed" $ TestCase $ do
        let input =
                "station,value\n"
                    <> "Alpha,1\nBeta,2\nAlpha,3\nBeta,4\n"
                    <> "Alpha,5\nBeta,6\nAlpha,7\nBeta,8\n"
            schema = Schema $ M.fromList [("station", schemaType @T.Text)]
        withRawCsv input $ \path -> do
            df <-
                L.runDataFrame
                    (L.scanCsvStreaming schema (T.pack path))
            case D.unsafeGetColumn "station" df of
                PackedText{} -> pure ()
                _ -> assertFailure "expected the CSV text column to stay packed"

tests :: Test
tests =
    TestList
        [ streamingPathReaderFillsWindow
        , streamingWindowBoundaryIsLossless
        , streamingQuotedNewlineSpansWindow
        , streamingDefaultMatchesEager
        , streamingReaderFailurePropagates
        , streamingKeepsTextPacked
        ]

main :: IO ()
main = do
    result <- runTestTT tests
    when (failures result > 0 || errors result > 0) Exit.exitFailure