diff --git a/dataframe-core/src-internal/DataFrame/Internal/Column/Operations.hs b/dataframe-core/src-internal/DataFrame/Internal/Column/Operations.hs index 750ac9a7..76913cb5 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Column/Operations.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Column/Operations.hs @@ -248,7 +248,23 @@ sliceColumn start n (BoxedColumn bm xs) = BoxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs) sliceColumn start n (UnboxedColumn bm xs) = UnboxedColumn (fmap (bitmapSlice start n) bm) (VG.slice start n xs) -sliceColumn start n c@(PackedText _ _) = sliceColumn start n (materializePacked c) +sliceColumn start n (PackedText bm p) + -- packedGather decodes an out-of-range index as the empty string, where + -- the boxed and unboxed arms reject the slice, so check first. + | start < 0 || n < 0 || start + n > packedLength p = + errorWithoutStackTrace + ( "sliceColumn: invalid slice (" + ++ show start + ++ "," + ++ show n + ++ "," + ++ show (packedLength p) + ++ ")" + ) + | otherwise = + PackedText + (fmap (bitmapSlice start n) bm) + (packedGather (VU.enumFromN start n) p) {-# INLINE sliceColumn #-} -- | O(n) Selects the elements at a given set of indices. Does not change the order. diff --git a/tests/Internal/PackedText.hs b/tests/Internal/PackedText.hs index e85c5637..2bc7a01e 100644 --- a/tests/Internal/PackedText.hs +++ b/tests/Internal/PackedText.hs @@ -14,6 +14,7 @@ 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) @@ -169,6 +170,40 @@ leftJoinSentinelPreservesPacked = TestCase $ do (DI.isPackedText (unsafeGetColumn "v" joinedP)) assertBool "left join packed == boxed" (joinedP == joinedB) +-- A slice of a packed column stays packed: decoding to Text would make a +-- ten-row slice cost the whole frame. +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) + +-- Every other representation rejects a slice running past the end; the packed +-- arm must not quietly decode the overrun as empty strings. +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) + tests :: [Test] tests = [ TestLabel "PackedText display parity" displayParity @@ -183,4 +218,7 @@ tests = , 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 ]