From a132d1b269ebece77eac0bd2540de22191ead7bc Mon Sep 17 00:00:00 2001 From: Samuel Schlesinger Date: Thu, 20 Aug 2026 18:47:40 -0400 Subject: [PATCH 1/2] test: Pin grouped aggregation over nullable columns sliceGroups permutes a column's data into group order but slices the null bitmap at the group's offsets in the original row order, so a group's values and its validity bits come from different rows unless the frame already happens to be sorted by the key. Both cases use interleaved keys so valueIndices is not the identity permutation. Currently the first reports a=4.0 for a group summing to 10.0; the second mis-attributes nulls in both groups. --- tests/Operations/Nullable.hs | 72 +++++++++++++++++++++++++++++++++++- 1 file changed, 70 insertions(+), 2 deletions(-) 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 From 449877422881191f2501c83d119330f84eb4c711 Mon Sep 17 00:00:00 2001 From: Samuel Schlesinger Date: Thu, 20 Aug 2026 18:47:40 -0400 Subject: [PATCH 2/2] fix: Permute the group bitmap with the data in sliceGroups The data was backpermuted into group order while the bitmap was sliced at the group's offsets in the original row order, so values and validity bits came from different rows unless the frame was already sorted by the key. Permute the bitmap by the same indices, once, before slicing. --- .../src-internal/DataFrame/Internal/Interpreter.hs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) 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