diff --git a/dataframe-core/src-internal/DataFrame/Internal/Column/Bitmap.hs b/dataframe-core/src-internal/DataFrame/Internal/Column/Bitmap.hs index da9cd6d9..dd381327 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Column/Bitmap.hs +++ b/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 = diff --git a/dataframe-core/src-internal/DataFrame/Internal/Column/Properties.hs b/dataframe-core/src-internal/DataFrame/Internal/Column/Properties.hs index cf06c1b6..c7dc9810 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Column/Properties.hs +++ b/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 #-} diff --git a/tests/Operations/Statistics.hs b/tests/Operations/Statistics.hs index db907654..4be37299 100644 --- a/tests/Operations/Statistics.hs +++ b/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