diff --git a/dataframe-core/src-internal/DataFrame/Internal/Interpreter.hs b/dataframe-core/src-internal/DataFrame/Internal/Interpreter.hs index 414352cc..34383dc8 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Interpreter.hs +++ b/dataframe-core/src-internal/DataFrame/Internal/Interpreter.hs @@ -477,20 +477,27 @@ sliceGroups col os indices = case col of V.generate (VU.length indices) ((vec `V.unsafeIndex`) . (indices `VU.unsafeIndex`)) + !sortedBm = permuteBitmap bm in V.generate nGroups $ \i -> BoxedColumn - (fmap (bitmapSlice (start i) (len i)) bm) + (fmap (bitmapSlice (start i) (len i)) sortedBm) (V.unsafeSlice (start i) (len i) sorted) UnboxedColumn bm vec -> let !sorted = VU.unsafeBackpermute vec indices + !sortedBm = permuteBitmap bm in V.generate nGroups $ \i -> UnboxedColumn - (fmap (bitmapSlice (start i) (len i)) bm) + (fmap (bitmapSlice (start i) (len i)) sortedBm) (VU.unsafeSlice (start i) (len i) sorted) where !nGroups = VU.length os - 1 start i = os `VU.unsafeIndex` i len i = os `VU.unsafeIndex` (i + 1) - start i + permuteBitmap = fmap $ \bm -> + buildBitmapFromValid $ + VU.map + (\r -> if bitmapTestBit bm r then 1 else 0) + indices {-# INLINE sliceGroups #-} numGroups :: GroupedDataFrame -> Int diff --git a/tests/Operations/Nullable.hs b/tests/Operations/Nullable.hs index 0bcac6bc..36f6d403 100644 --- a/tests/Operations/Nullable.hs +++ b/tests/Operations/Nullable.hs @@ -5,9 +5,11 @@ module Operations.Nullable where +import Data.Function ((&)) +import qualified Data.Text as T import qualified Data.Vector as V import qualified DataFrame as D -import DataFrame.Expression.Operators ((.*), (.+), (.-), (./), (.==)) +import DataFrame.Expression.Operators (as, (.*), (.+), (.-), (./), (.==)) import qualified DataFrame.Functions as F import qualified DataFrame.Internal.Column as DI import qualified DataFrame.Internal.DataFrame as DI @@ -723,9 +725,75 @@ applyPlainInt = ) ) +-- --------------------------------------------------------------------------- +-- Grouped aggregation over nullable columns +-- --------------------------------------------------------------------------- + +-- | Interleaved, so grouping's row permutation is not the identity. +interleavedKeys :: [T.Text] +interleavedKeys = ["a", "b", "a", "b", "a", "b", "a", "b"] + +-- | Group @a@ holds 1,2,3,4; group @b@ is entirely null. +nullsInOneGroupValues :: [Maybe Double] +nullsInOneGroupValues = + [Just 1, Nothing, Just 2, Nothing, Just 3, Nothing, Just 4, Nothing] + +nullsInOneGroup :: D.DataFrame +nullsInOneGroup = + D.fromNamedColumns + [ ("k", DI.fromList interleavedKeys) + , ("v", DI.fromVector (V.fromList nullsInOneGroupValues)) + ] + +-- | Group @a@ sums to 7, group @b@ to 40. +nullsInBothGroupsValues :: [Maybe Double] +nullsInBothGroupsValues = + [Just 1, Just 10, Just 2, Nothing, Nothing, Just 30, Just 4, Nothing] + +nullsInBothGroups :: D.DataFrame +nullsInBothGroups = + D.fromNamedColumns + [ ("k", DI.fromList interleavedKeys) + , ("v", DI.fromVector (V.fromList nullsInBothGroupsValues)) + ] + +sumGroupedNullable :: D.DataFrame -> D.DataFrame +sumGroupedNullable df = + df + & D.groupBy ["k"] + & D.aggregate [F.sumMaybe (F.col @(Maybe Double) "v") `as` "s"] + & D.sortBy [D.Asc (F.col @T.Text "k")] + +expectedGroupSums :: [Double] -> D.DataFrame +expectedGroupSums sums = + D.fromNamedColumns + [ ("k", DI.fromList (["a", "b"] :: [T.Text])) + , ("s", DI.fromList sums) + ] + +sumMaybeOverInterleavedGroups :: Test +sumMaybeOverInterleavedGroups = + TestCase + ( assertEqual + "sumMaybe sums each group's own non-null values" + (expectedGroupSums [10.0, 0.0]) + (sumGroupedNullable nullsInOneGroup) + ) + +sumMaybeWithNullsInEveryGroup :: Test +sumMaybeWithNullsInEveryGroup = + TestCase + ( assertEqual + "sumMaybe skips nulls in every group, not just one" + (expectedGroupSums [7.0, 40.0]) + (sumGroupedNullable nullsInBothGroups) + ) + tests :: [Test] tests = - [ TestLabel "addIntMaybeInt" addIntMaybeInt + [ TestLabel "sumMaybeOverInterleavedGroups" sumMaybeOverInterleavedGroups + , TestLabel "sumMaybeWithNullsInEveryGroup" sumMaybeWithNullsInEveryGroup + , TestLabel "addIntMaybeInt" addIntMaybeInt , TestLabel "addMaybeIntInt" addMaybeIntInt , TestLabel "addIntInt" addIntInt , TestLabel "addMaybeMaybe" addMaybeMaybe