PackedText.hs

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

-- | Correctness oracle: packed-text columns behave like boxed Text columns.
module Internal.PackedText (tests) where

import qualified Data.Text as T
import qualified Data.Text.Array as A
import qualified Data.Vector.Unboxed as VU

import Control.Exception (SomeException, evaluate, try)
import Control.Monad (zipWithM_)
import qualified Data.ByteString as B
import Data.Text.Encoding (encodeUtf8)
import Data.Word (Word8)
import qualified DataFrame as D
import DataFrame.Expression.Operators ((./=.), (.==.))
import qualified DataFrame.Functions as F
import qualified DataFrame.Internal.Column as DI
import DataFrame.Internal.Data.PackedText (mkPackedContiguous)
import DataFrame.Internal.DataFrame (unsafeGetColumn)
import qualified DataFrame.Operations.Aggregation as Agg
import Test.HUnit

-- Sample string corpus, including empty, ASCII, and multibyte UTF-8 fields.
sampleRows :: [T.Text]
sampleRows =
    [ "apple"
    , "banana"
    , ""
    , "apple"
    , "Zebra"
    , "café"
    , "naïve"
    , "日本語"
    , "banana"
    , "apple"
    ]

-- Build an A.Array from a flat byte list.
arrayFromBytes :: [Word8] -> A.Array
arrayFromBytes ws = A.run $ do
    m <- A.new (length ws)
    zipWithM_ (A.unsafeWrite m) [0 ..] ws
    pure m

-- Build a PackedText column directly from raw UTF-8 bytes + offsets.
packedFromTexts :: [T.Text] -> DI.Column
packedFromTexts ts =
    let bytess = map (B.unpack . encodeUtf8) ts
        flat = concat bytess
        offs = scanl (+) 0 (map length bytess)
        arr = arrayFromBytes flat
     in DI.PackedText Nothing (mkPackedContiguous arr (VU.fromList offs))

boxedFromTexts :: [T.Text] -> DI.Column
boxedFromTexts = DI.fromList

packedCol, boxedCol :: DI.Column
packedCol = packedFromTexts sampleRows
boxedCol = boxedFromTexts sampleRows

-- Single-column dataframes for groupBy / hashing.
packedDf, boxedDf :: D.DataFrame
packedDf = D.fromColumns [("k", packedCol)]
boxedDf = D.fromColumns [("k", boxedCol)]

displayParity :: Test
displayParity =
    TestCase $
        assertEqual "show packed == show boxed" (show boxedCol) (show packedCol)

columnEqParity :: Test
columnEqParity = TestCase $ do
    assertBool "packed == boxed" (packedCol == boxedCol)
    assertBool "boxed == packed" (boxedCol == packedCol)
    assertBool "packed == packed" (packedCol == packedFromTexts sampleRows)

extractionParity :: Test
extractionParity =
    TestCase $
        assertEqual
            "toList @Text packed == boxed"
            (DI.toList @T.Text boxedCol)
            (DI.toList @T.Text packedCol)

orderingParity :: Test
orderingParity = TestCase $ do
    let sortedPacked = D.sortBy [D.Asc (F.col @T.Text "k")] packedDf
        sortedBoxed = D.sortBy [D.Asc (F.col @T.Text "k")] boxedDf
    assertBool "asc sort parity" (sortedPacked == sortedBoxed)
    let descPacked = D.sortBy [D.Desc (F.col @T.Text "k")] packedDf
        descBoxed = D.sortBy [D.Desc (F.col @T.Text "k")] boxedDf
    assertBool "desc sort parity" (descPacked == descBoxed)

hashingParity :: Test
hashingParity =
    TestCase $
        assertEqual
            "row hashes equal"
            (Agg.computeRowHashes [0] boxedDf)
            (Agg.computeRowHashes [0] packedDf)

groupByParity :: Test
groupByParity =
    TestCase $
        assertEqual
            "groupBy + distinct parity"
            (D.distinct boxedDf)
            (D.distinct packedDf)

joinParity :: Test
joinParity = TestCase $ do
    let lhsP = D.fromColumns [("k", packedCol)]
        lhsB = D.fromColumns [("k", boxedCol)]
        rhs = D.fromColumns [("k", boxedFromTexts ["apple", "banana", "日本語"])]
        joinedP = D.innerJoin ["k"] lhsP rhs
        joinedB = D.innerJoin ["k"] lhsB rhs
    assertBool "inner join parity" (joinedP == joinedB)

-- A gather on a packed column stays packed and equals the boxed gather.
gatherPreservesPacked :: Test
gatherPreservesPacked = TestCase $ do
    let ixs = VU.fromList [9, 0, 0, 5, 7, 2, 3 :: Int]
        gp = DI.atIndicesStable ixs packedCol
        gb = DI.atIndicesStable ixs boxedCol
    assertBool "gathered packed stays PackedText" (DI.isPackedText gp)
    assertBool "gathered packed == gathered boxed" (gp == gb)
    assertEqual
        "gathered packed toList == boxed"
        (DI.toList @T.Text gb)
        (DI.toList @T.Text gp)

-- A sort on a packed column stays packed and equals the boxed sort.
sortPreservesPacked :: Test
sortPreservesPacked = TestCase $ do
    let sortedP = D.sortBy [D.Asc (F.col @T.Text "k")] packedDf
        sortedB = D.sortBy [D.Asc (F.col @T.Text "k")] boxedDf
    assertBool
        "sorted packed column stays PackedText"
        (DI.isPackedText (unsafeGetColumn "k" sortedP))
    assertBool "sorted packed == sorted boxed" (sortedP == sortedB)

-- A left join on a packed right column stays packed and equals boxed.
leftJoinSentinelPreservesPacked :: Test
leftJoinSentinelPreservesPacked = TestCase $ do
    let lhsP = D.fromColumns [("k", packedCol)]
        lhsB = D.fromColumns [("k", boxedCol)]
        rhsP =
            D.fromColumns
                [ ("k", packedFromTexts ["apple", "banana", "日本語"])
                , ("v", packedFromTexts ["RA", "RB", "RC"])
                ]
        rhsB =
            D.fromColumns
                [ ("k", boxedFromTexts ["apple", "banana", "日本語"])
                , ("v", boxedFromTexts ["RA", "RB", "RC"])
                ]
        joinedP = D.leftJoin ["k"] lhsP rhsP
        joinedB = D.leftJoin ["k"] lhsB rhsB
    assertBool
        "left-joined packed right column stays PackedText"
        (DI.isPackedText (unsafeGetColumn "v" joinedP))
    assertBool "left join packed == boxed" (joinedP == joinedB)

-- A slice of a packed column stays packed instead of decoding to Text.
slicePreservesPacked :: Test
slicePreservesPacked = TestCase $ do
    let sp = DI.sliceColumn 3 4 packedCol
        sb = DI.sliceColumn 3 4 boxedCol
    assertBool "sliced packed stays PackedText" (DI.isPackedText sp)
    assertBool "sliced packed == sliced boxed" (sp == sb)
    assertEqual
        "sliced packed toList == boxed"
        (DI.toList @T.Text sb)
        (DI.toList @T.Text sp)

-- takeLastColumn is sliceColumn from an offset, so it must stay packed too.
takeLastPreservesPacked :: Test
takeLastPreservesPacked = TestCase $ do
    let tp = DI.takeLastColumn 3 packedCol
        tb = DI.takeLastColumn 3 boxedCol
    assertBool "takeLast packed stays PackedText" (DI.isPackedText tp)
    assertEqual
        "takeLast packed toList == boxed"
        (DI.toList @T.Text tb)
        (DI.toList @T.Text tp)

-- A packed slice past the end must be rejected, not decoded as empty.
sliceRejectsInvalidBounds :: Test
sliceRejectsInvalidBounds = TestCase $ do
    r <- try (evaluate (DI.columnLength (DI.sliceColumn 8 5 packedCol)))
    case r :: Either SomeException Int of
        Left _ -> pure ()
        Right len ->
            assertFailure ("expected an invalid slice, got length " ++ show len)

-- Rows for the eq/neq kernel: valid text, empty, and invalid UTF-8.
kernelRowBytes :: [[Word8]]
kernelRowBytes =
    map (B.unpack . encodeUtf8) ["apple", "", "café", "日本語", "apple", "\xFFFD"]
        ++ [[0xFF], [0x63, 0x61, 0x66, 0xC3], [0x80], [0xEF, 0xBF]]

packedFromBytes :: [[Word8]] -> DI.Column
packedFromBytes rows =
    DI.PackedText
        Nothing
        ( mkPackedContiguous
            (arrayFromBytes (concat rows))
            (VU.fromList (scanl (+) 0 (map length rows)))
        )

-- The eq/neq-literal kernel must agree with the decoded-column expression.
eqLiteralKernelParity :: Test
eqLiteralKernelParity = TestCase $ do
    let base = D.fromColumns [("k", packedFromBytes kernelRowBytes)]
        sorted = D.sortBy [D.Asc (F.col @T.Text "k")] base
        boxed df = D.fromColumns [("k", DI.materializePacked (unsafeGetColumn "k" df))]
        lits = ["apple", "", "café", "日本語", "\xFFFD", "caf", "missing"] :: [T.Text]
        result :: D.DataFrame -> D.Expr Bool -> [Bool]
        result df e = DI.toList @Bool (unsafeGetColumn "r" (D.derive "r" e df))
    assertBool
        "sorted payload stays PackedText"
        (DI.isPackedText (unsafeGetColumn "k" sorted))
    sequence_
        [ assertEqual
            (label ++ " " ++ show lit)
            (result (boxed df) e)
            (result df e)
        | (label, df) <- [("base", base), ("sorted", sorted)]
        , lit <- lits
        , e <-
            [ F.col @T.Text "k" .==. F.lit lit
            , F.col @T.Text "k" ./=. F.lit lit
            , F.lit lit .==. F.col @T.Text "k"
            , F.lit lit ./=. F.col @T.Text "k"
            ]
        ]

tests :: [Test]
tests =
    [ TestLabel "PackedText display parity" displayParity
    , TestLabel "PackedText equality parity" columnEqParity
    , TestLabel "PackedText extraction parity" extractionParity
    , TestLabel "PackedText ordering parity" orderingParity
    , TestLabel "PackedText hashing parity" hashingParity
    , TestLabel "PackedText groupBy parity" groupByParity
    , TestLabel "PackedText join parity" joinParity
    , TestLabel "PackedText gather preserves packed" gatherPreservesPacked
    , TestLabel "PackedText sort preserves packed" sortPreservesPacked
    , TestLabel
        "PackedText left-join sentinel preserves packed"
        leftJoinSentinelPreservesPacked
    , TestLabel "PackedText slice preserves packed" slicePreservesPacked
    , TestLabel "PackedText takeLast preserves packed" takeLastPreservesPacked
    , TestLabel "PackedText slice rejects bad bounds" sliceRejectsInvalidBounds
    , TestLabel "PackedText eq/neq literal kernel parity" eqLiteralKernelParity
    ]