fix: Count validity bits only up to the column length

dataframe-core/src-internal/DataFrame/Internal/Column/Bitmap.hs

@@ -4,7 +4,16 @@ module DataFrame.Internal.Column.Bitmap where
 
 import Control.Monad (foldM_, forM_, when)
 import Control.Monad.ST (ST, runST)
-import Data.Bits (complement, setBit, shiftL, shiftR, testBit, (.&.), (.|.))
+import Data.Bits (
+    complement,
+    popCount,
+    setBit,
+    shiftL,
+    shiftR,
+    testBit,
+    (.&.),
+    (.|.),
+ )
 import Data.List (foldl')
 import Data.Maybe (fromMaybe, isNothing)
 import qualified Data.Vector.Unboxed as VU
@@ -100,6 +109,23 @@ bitmapSlice start len bm
                 VU.generate n $
                     \i -> if bitmapTestBit bm (start + i) then 1 else 0
 
+{- | Count the set bits among the first @n@ bits of a bitmap. A bitmap does
+not know the length of the column it describes, and 'bitmapSlice' keeps whole
+bytes on its aligned path, so the bits past @n@ may still describe rows
+outside the slice.
+-}
+popCountUpTo :: Int -> Bitmap -> Int
+popCountUpTo n bm = whole + partial
+  where
+    !fullBytes = min (n `shiftR` 3) (VU.length bm)
+    !rest = n .&. 7
+    whole = VU.foldl' (\acc b -> acc + popCount b) 0 (VU.take fullBytes bm)
+    partial
+        | rest == 0 || fullBytes >= VU.length bm = 0
+        | otherwise =
+            popCount (VU.unsafeIndex bm fullBytes .&. ((1 `shiftL` rest) - 1))
+{-# INLINE popCountUpTo #-}
+
 -- | Concatenate two bitmaps covering @n1@ and @n2@ rows respectively.
 bitmapConcat :: Int -> Bitmap -> Int -> Bitmap -> Bitmap
 bitmapConcat n1 bm1 n2 bm2 =

dataframe-core/src-internal/DataFrame/Internal/Column/Properties.hs

@@ -25,7 +25,6 @@ import qualified Data.Text as T
 import qualified Data.Vector as VB
 import qualified Data.Vector.Unboxed as VU
 
-import Data.Bits (popCount)
 import Data.Kind (Type)
 import Data.Maybe (isJust)
 import Data.Type.Equality (TestEquality (..))
@@ -64,9 +63,12 @@ hasMissing _ = False
 
 -- | Checks if a column contains only missing values.
 allMissing :: Column -> Bool
-allMissing (BoxedColumn (Just bm) col) = VU.all (== 0) bm && not (VB.null col)
-allMissing (UnboxedColumn (Just bm) col) = VU.all (== 0) bm && not (VU.null col)
-allMissing (PackedText (Just bm) p) = VU.all (== 0) bm && packedLength p > 0
+allMissing (BoxedColumn (Just bm) col) =
+    not (VB.null col) && popCountUpTo (VB.length col) bm == 0
+allMissing (UnboxedColumn (Just bm) col) =
+    not (VU.null col) && popCountUpTo (VU.length col) bm == 0
+allMissing (PackedText (Just bm) p) =
+    packedLength p > 0 && popCountUpTo (packedLength p) bm == 0
 allMissing _ = False
 
 -- | Checks if a column contains numeric values.
@@ -143,9 +145,9 @@ columnElemIsNull _ _ = False
 numElements :: Column -> Int
 numElements (MergedColumn a b) = min (columnLength a) (columnLength b)
 numElements (BoxedColumn Nothing xs) = VB.length xs
-numElements (BoxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
+numElements (BoxedColumn (Just bm) xs) = popCountUpTo (VB.length xs) bm
 numElements (UnboxedColumn Nothing xs) = VU.length xs
-numElements (UnboxedColumn (Just bm) _xs) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
+numElements (UnboxedColumn (Just bm) xs) = popCountUpTo (VU.length xs) bm
 numElements (PackedText Nothing p) = packedLength p
-numElements (PackedText (Just bm) _p) = VU.foldl' (\acc b -> acc + popCount b) 0 bm
+numElements (PackedText (Just bm) p) = popCountUpTo (packedLength p) bm
 {-# INLINE numElements #-}

tests/Operations/Statistics.hs

@@ -7,6 +7,7 @@ module Operations.Statistics where
 import qualified Data.Vector.Unboxed as VU
 import qualified DataFrame as D
 import qualified DataFrame.Internal.Column as DI
+import DataFrame.Internal.DataFrame (getColumn)
 import qualified DataFrame.Internal.Statistics as D
 
 import Assertions
@@ -198,6 +199,56 @@ summarizeOptional =
             )
         )
 
+{- | A sliced column's bitmap keeps whole bytes, so its trailing bits still
+describe rows past the end of the slice. Counting non-null rows has to stop
+at the column's length rather than fold the whole byte vector.
+-}
+sixteenNullableRows :: D.DataFrame
+sixteenNullableRows =
+    D.fromNamedColumns
+        [ ("x", D.fromList (map (Just . fromIntegral) [1 .. 16 :: Int] :: [Maybe Double]))
+        ]
+
+countAfterTake :: Test
+countAfterTake =
+    TestCase
+        ( assertEqual
+            "summarize counts the rows a take actually kept"
+            (Just (DI.fromList ([3.0] :: [Double])))
+            ( getColumn "x" $
+                D.take 1 $
+                    D.summarize (D.take 3 sixteenNullableRows)
+            )
+        )
+
+countAfterRange :: Test
+countAfterRange =
+    TestCase
+        ( assertEqual
+            "summarize counts the rows a range actually kept"
+            (Just (DI.fromList ([5.0] :: [Double])))
+            ( getColumn "x" $
+                D.take 1 $
+                    D.summarize (D.range (8, 13) sixteenNullableRows)
+            )
+        )
+
+-- allMissing folds the same bitmap, so it needs the same length cutoff: the
+-- padding bits of a sliced all-null column would otherwise read as present.
+allMissingAfterSlice :: Test
+allMissingAfterSlice =
+    TestCase
+        ( assertEqual
+            "an all-null slice is still all-null"
+            (Just True)
+            (DI.allMissing <$> getColumn "x" (D.take 3 allNullRows))
+        )
+
+allNullRows :: D.DataFrame
+allNullRows =
+    D.fromNamedColumns
+        [("x", DI.fromList (replicate 16 (Nothing :: Maybe Double)))]
+
 -- correlation
 
 correlationDf :: D.DataFrame
@@ -277,6 +328,9 @@ tests =
     , TestLabel "wrongQuantileNumber" wrongQuantileNumber
     , TestLabel "wrongQuantileIndex" wrongQuantileIndex
     , TestLabel "summarizeOptional" summarizeOptional
+    , TestLabel "countAfterTake" countAfterTake
+    , TestLabel "countAfterRange" countAfterRange
+    , TestLabel "allMissingAfterSlice" allMissingAfterSlice
     , TestLabel "correlationPerfectPositive" correlationPerfectPositive
     , TestLabel "correlationPerfectNegative" correlationPerfectNegative
     , TestLabel "correlationSelfIdentity" correlationSelfIdentity