fix: frame invariants (sortBy, selectRows, range, rename)

dataframe-core/src-internal/DataFrame/Errors.hs

@@ -32,9 +32,11 @@ data DataFrameException where
         DataFrameException
     AggregatedAndNonAggregatedException :: T.Text -> T.Text -> DataFrameException
     ColumnsNotFoundException :: [T.Text] -> T.Text -> [T.Text] -> DataFrameException
+    DuplicateColumnException :: T.Text -> T.Text -> DataFrameException
     EmptyDataSetException :: T.Text -> DataFrameException
     InternalException :: T.Text -> DataFrameException
     NonColumnReferenceException :: T.Text -> DataFrameException
+    RowsOutOfBoundsException :: [Int] -> Int -> DataFrameException
     UnaggregatedException :: T.Text -> DataFrameException
     WrongQuantileNumberException :: Int -> DataFrameException
     WrongQuantileIndexException :: VU.Vector Int -> Int -> DataFrameException
@@ -54,7 +56,20 @@ instance Show DataFrameException where
                 (callingFunctionName context)
                 errorString
     show (ColumnsNotFoundException columnNames callPoint availableColumns) = columnsNotFound columnNames callPoint availableColumns
+    show (DuplicateColumnException name callPoint) =
+        red "\n\n[ERROR] "
+            ++ "Column already exists: "
+            ++ T.unpack name
+            ++ " for operation "
+            ++ T.unpack callPoint
     show (EmptyDataSetException callPoint) = emptyDataSetError callPoint
+    show (RowsOutOfBoundsException ixs n) =
+        red "\n\n[ERROR] "
+            ++ "Row indexes out of bounds: "
+            ++ show ixs
+            ++ " (the dataframe has "
+            ++ show n
+            ++ " rows)"
     show (WrongQuantileNumberException q) = wrongQuantileNumberError q
     show (WrongQuantileIndexException qs q) = wrongQuantileIndexError qs q
     show (InternalException msg) = "Internal error: " ++ T.unpack msg

dataframe-operations/src/DataFrame/Operations/Core.hs

@@ -496,6 +496,9 @@ renameSafe ::
     T.Text -> T.Text -> DataFrame -> Either DataFrameException DataFrame
 renameSafe orig new df
     | null df = throw (EmptyDataSetException "rename")
+    -- Renaming onto a live column would orphan it and corrupt the frame.
+    | orig /= new && M.member new (columnIndices df) =
+        Left (DuplicateColumnException new "rename")
     | otherwise = fromMaybe
         (Left $ ColumnsNotFoundException [orig] "rename" (M.keys $ columnIndices df))
         $ do

dataframe-operations/src/DataFrame/Operations/Permutation.hs

@@ -27,7 +27,7 @@ import Control.Exception (throw)
 import Control.Monad.ST (runST)
 import Data.Type.Equality (testEquality, (:~:) (Refl))
 import Data.Vector.Internal.Check (HasCallStack)
-import DataFrame.Errors (DataFrameException (..))
+import DataFrame.Errors (DataFrameException (..), TypeErrorContext (..))
 import DataFrame.Internal.Column (
     Column (..),
     Columnable,
@@ -44,7 +44,7 @@ import DataFrame.Internal.PackedText (packedSlice, sliceCmpBytes)
 import DataFrame.Operations.Core (dimensions)
 import DataFrame.Operations.Transformations (derive)
 import System.Random (Random (randomR), RandomGen)
-import Type.Reflection (typeRep)
+import Type.Reflection (TypeRep, Typeable, typeRep)
 
 -- | Sort order taken as a parameter by the 'sortBy' function.
 data SortOrder where
@@ -127,42 +127,56 @@ sortOrderComparator (Asc (Col name :: Expr a)) df =
     case unsafeGetColumn name df of
         BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
             Just Refl -> \i j -> compare (v `V.unsafeIndex` i) (v `V.unsafeIndex` j)
-            Nothing -> \_ _ -> EQ
+            Nothing -> sortTypeMismatch name (typeRep @a) (typeRep @b)
         UnboxedColumn _ (v :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
             Just Refl -> \i j -> compare (v `VU.unsafeIndex` i) (v `VU.unsafeIndex` j)
-            Nothing -> \_ _ -> EQ
+            Nothing -> sortTypeMismatch name (typeRep @a) (typeRep @b)
         PackedText _ p -> case testEquality (typeRep @a) (typeRep @T.Text) of
             Just Refl -> \i j ->
                 let (ai, oi, li) = packedSlice p i
                     (aj, oj, lj) = packedSlice p j
                  in sliceCmpBytes ai oi li aj oj lj
-            Nothing -> \_ _ -> EQ
+            Nothing -> sortTypeMismatch name (typeRep @a) (typeRep @T.Text)
         c@(MergedColumn _ _) -> case materializeMerged c of
             BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
                 Just Refl -> \i j -> compare (v `V.unsafeIndex` i) (v `V.unsafeIndex` j)
-                Nothing -> \_ _ -> EQ
-            _ -> \_ _ -> EQ
+                Nothing -> sortTypeMismatch name (typeRep @a) (typeRep @b)
+            _ -> throw (InternalException "sortBy: unsupported merged column")
 sortOrderComparator (Desc (Col name :: Expr a)) df =
     case unsafeGetColumn name df of
         BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
             Just Refl -> \i j -> compare (v `V.unsafeIndex` j) (v `V.unsafeIndex` i)
-            Nothing -> \_ _ -> EQ
+            Nothing -> sortTypeMismatch name (typeRep @a) (typeRep @b)
         UnboxedColumn _ (v :: VU.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
             Just Refl -> \i j -> compare (v `VU.unsafeIndex` j) (v `VU.unsafeIndex` i)
-            Nothing -> \_ _ -> EQ
+            Nothing -> sortTypeMismatch name (typeRep @a) (typeRep @b)
         PackedText _ p -> case testEquality (typeRep @a) (typeRep @T.Text) of
             Just Refl -> \i j ->
                 let (ai, oi, li) = packedSlice p i
                     (aj, oj, lj) = packedSlice p j
                  in sliceCmpBytes aj oj lj ai oi li
-            Nothing -> \_ _ -> EQ
+            Nothing -> sortTypeMismatch name (typeRep @a) (typeRep @T.Text)
         c@(MergedColumn _ _) -> case materializeMerged c of
             BoxedColumn _ (v :: V.Vector b) -> case testEquality (typeRep @a) (typeRep @b) of
                 Just Refl -> \i j -> compare (v `V.unsafeIndex` j) (v `V.unsafeIndex` i)
-                Nothing -> \_ _ -> EQ
-            _ -> \_ _ -> EQ
+                Nothing -> sortTypeMismatch name (typeRep @a) (typeRep @b)
+            _ -> throw (InternalException "sortBy: unsupported merged column")
 sortOrderComparator _ _ = error "Sorting on compound column"
 
+-- A wrong runtime type must fail, not compare every row as EQ.
+sortTypeMismatch ::
+    (Typeable a, Typeable b) => T.Text -> TypeRep a -> TypeRep b -> c
+sortTypeMismatch name userT colT =
+    throw $
+        TypeMismatchException
+            ( MkTypeErrorContext
+                { userType = Right userT
+                , expectedType = Right colT
+                , errorColumnName = Just (T.unpack name)
+                , callingFunctionName = Just "sortBy"
+                }
+            )
+
 -- | Sort row indices using a comparator function.
 sortIndices :: (Int -> Int -> Ordering) -> Int -> VU.Vector Int
 sortIndices cmp nRows = runST $ do

dataframe-operations/src/DataFrame/Operations/Subset.hs

@@ -148,12 +148,14 @@ dropLast n d =
 range :: (Int, Int) -> DataFrame -> DataFrame
 range (start, end) d =
     d
-        { columns = V.map (sliceColumn (clip start 0 r) n') (columns d)
+        { columns = V.map (sliceColumn start' n') (columns d)
         , dataframeDimensions = (n', c)
         }
   where
     (r, c) = dataframeDimensions d
-    n' = clip (end - start) 0 r
+    start' = clip start 0 r
+    -- Bounded by the rows left after start', not by r.
+    n' = clip (clip end 0 r - start') 0 (r - start')
 
 clip :: Int -> Int -> Int -> Int
 clip n left right = min right $ max n left
@@ -424,12 +426,17 @@ selectBy xs df = select finalSelection df
 > selectRows [0, 2, 4] df
 -}
 selectRows :: [Int] -> DataFrame -> DataFrame
-selectRows ixs df =
-    df
-        { columns = V.map (atIndicesStable ixs') (columns df)
-        , dataframeDimensions = (VU.length ixs', snd (dataframeDimensions df))
-        }
+selectRows ixs df
+    -- 'atIndicesStable' gathers with unsafeIndex; bounds-check here.
+    | not (L.null oob) = throw (RowsOutOfBoundsException oob r)
+    | otherwise =
+        df
+            { columns = V.map (atIndicesStable ixs') (columns df)
+            , dataframeDimensions = (VU.length ixs', snd (dataframeDimensions df))
+            }
   where
+    (r, _) = dataframeDimensions df
+    oob = L.filter (\i -> i < 0 || i >= r) ixs
     ixs' = VU.fromList ixs
 
 {- | O(n) inverse of select

tests/Operations/Core.hs

@@ -105,9 +105,29 @@ fromRowsRoundTripsWithNulls =
                 (D.fromRows (D.columnNames df) (map (map snd) (D.toRowList df)))
         )
 
+renameOntoExistingColumn :: Test
+renameOntoExistingColumn =
+    TestCase
+        ( assertExpectException
+            "[Error Case]"
+            "Column already exists: B"
+            (print $ D.rename "A" "B" testData)
+        )
+
+renameToItselfIsIdentity :: Test
+renameToItselfIsIdentity =
+    TestCase
+        ( assertEqual
+            "renaming a column to itself is a no-op"
+            testData
+            (D.rename "A" "A" testData)
+        )
+
 tests :: [Test]
 tests =
     [ TestLabel "createsDataFrameFromRows" createsDataFrameFromRows
+    , TestLabel "renameOntoExistingColumn" renameOntoExistingColumn
+    , TestLabel "renameToItselfIsIdentity" renameToItselfIsIdentity
     , TestLabel "fromRowsThrowsOnTypeMismatch" fromRowsThrowsOnTypeMismatch
     , TestLabel "fromRowsThrowsOnShortRow" fromRowsThrowsOnShortRow
     , TestLabel "fromRowsKeepsNullsInPlace" fromRowsKeepsNullsInPlace

tests/Operations/Sort.hs

@@ -12,6 +12,7 @@ import qualified DataFrame.Internal.Column as DI
 import System.Random
 import System.Random.Shuffle (shuffle')
 import Test.HUnit
+import Type.Reflection (typeRep)
 
 values :: [(T.Text, DI.Column)]
 values =
@@ -89,6 +90,15 @@ sortByColumnDoesNotExist =
             (print $ D.sortBy [D.Asc (F.col @Int "test0")] testData)
         )
 
+sortByWrongColumnType :: Test
+sortByWrongColumnType =
+    TestCase
+        ( assertExpectException
+            "[Error Case]"
+            (D.typeMismatchError (show $ typeRep @Double) (show $ typeRep @Int))
+            (print $ D.sortBy [D.Asc (F.col @Double "test1")] testData)
+        )
+
 compoundTestData :: D.DataFrame
 compoundTestData =
     D.fromNamedColumns
@@ -160,6 +170,7 @@ tests =
     [ TestLabel "sortByAscendingWAI" sortByAscendingWAI
     , TestLabel "sortByDescendingWAI" sortByDescendingWAI
     , TestLabel "sortByColumnDoesNotExist" sortByColumnDoesNotExist
+    , TestLabel "sortByWrongColumnType" sortByWrongColumnType
     , TestLabel "sortByTwoColumns" sortByTwoColumns
     , TestLabel "sortByOneColumnAscOneColumnDesc" sortByOneColumnAscOneColumnDesc
     , TestLabel "sortByCompoundExpression" sortByCompoundExpression

tests/Operations/Subset.hs

@@ -12,6 +12,8 @@ import DataFrame.Operations.Merge ()
 import System.Random
 import Test.HUnit
 
+import Assertions (assertExpectException)
+
 prop_dropZero :: DataFrame -> Bool
 prop_dropZero df = D.drop 0 df == df
 
@@ -53,6 +55,20 @@ prop_rangeFull df =
     let rows = fst (dataframeDimensions df)
      in D.range (0, rows) df == df
 
+prop_rangeClipsPastEnd :: DataFrame -> Bool
+prop_rangeClipsPastEnd df =
+    let rows = fst (dataframeDimensions df)
+     in D.range (rows `div` 2, rows + 10) df == D.drop (rows `div` 2) df
+
+prop_rangeClipsNegativeStart :: DataFrame -> Bool
+prop_rangeClipsNegativeStart df =
+    fst (dataframeDimensions (D.range (-5, 2) df))
+        == min 2 (fst (dataframeDimensions df))
+
+prop_selectRowsIdentity :: DataFrame -> Bool
+prop_selectRowsIdentity df =
+    D.selectRows [0 .. fst (dataframeDimensions df) - 1] df == df
+
 prop_selectAll :: DataFrame -> Bool
 prop_selectAll df = D.select (D.columnNames df) df == df
 
@@ -177,9 +193,31 @@ unit_stratifiedSplit_proportions =
                     )
                     (abs (vaProp - origProp) < tol)
 
+unit_selectRowsOutOfBounds :: Test
+unit_selectRowsOutOfBounds =
+    TestCase $ do
+        assertExpectException
+            "[Error Case]"
+            "Row indexes out of bounds"
+            (print $ D.selectRows [0, 1000000] strataDf)
+        assertExpectException
+            "[Error Case]"
+            "Row indexes out of bounds"
+            (print $ D.selectRows [-1] strataDf)
+
+unit_rangePastEndClips :: Test
+unit_rangePastEndClips =
+    TestCase $
+        assertEqual
+            "range past the end clips"
+            1
+            (fst (dataframeDimensions (D.range (9, 99) strataDf)))
+
 hunitTests :: [Test]
 hunitTests =
-    [ TestLabel "unit_stratifiedSample_full" unit_stratifiedSample_full
+    [ TestLabel "unit_selectRowsOutOfBounds" unit_selectRowsOutOfBounds
+    , TestLabel "unit_rangePastEndClips" unit_rangePastEndClips
+    , TestLabel "unit_stratifiedSample_full" unit_stratifiedSample_full
     , TestLabel "unit_stratifiedSplit_rowCount" unit_stratifiedSplit_rowCount
     , TestLabel
         "unit_stratifiedSplit_singleRowStratum"
@@ -199,6 +237,9 @@ tests =
     , prop_dropLastAll
     , prop_rangeEmpty
     , prop_rangeFull
+    , prop_rangeClipsPastEnd
+    , prop_rangeClipsNegativeStart
+    , prop_selectRowsIdentity
     , prop_selectAll
     , prop_selectEmpty
     , prop_excludeEmpty

tests/Properties/Categorical.hs

@@ -170,6 +170,16 @@ prop_renameRoundTrip (Frame df) =
              in notElem tmp (columnNames df) ==>
                     D.rename tmp name (D.rename name tmp df) === df
 
+-- | Renaming never changes the number of columns.
+prop_renameKeepsWidth :: Frame -> Property
+prop_renameKeepsWidth (Frame df) =
+    case columnNames df of
+        [] -> property True
+        (name : _) ->
+            let tmp = name <> "__rn_x"
+             in notElem tmp (columnNames df) ==>
+                    length (columnNames (D.rename name tmp df)) === D.nColumns df
+
 tests :: [Property]
 tests =
     [ property prop_unionCommutative
@@ -185,4 +195,5 @@ tests =
     , property prop_differenceDisjoint
     , property prop_excludeNothingIdentity
     , property prop_renameRoundTrip
+    , property prop_renameKeepsWidth
     ]