fix: Concatenate chunk bitmaps in a single pass

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

@@ -32,7 +32,7 @@ import Control.Monad (when)
 import Control.Monad.ST (runST)
 import Data.Bits (shiftR)
 import Data.Kind (Type)
-import Data.Maybe (fromMaybe, isNothing)
+import Data.Maybe (isNothing)
 import Data.Type.Equality (TestEquality (..))
 import DataFrame.Errors (
     DataFrameException (EmptyDataSetException, TypeMismatchException),
@@ -870,18 +870,7 @@ concatManyColumns (c0 : cs) = case c0 of
             rest = map getCol cs
             allVecs = v0 : map snd rest
             allBms = bm0 : map fst rest
-            newBm
-                | all isNothing allBms = Nothing
-                | otherwise =
-                    let pairs = zip allVecs allBms
-                        expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VB.length v)) mb) pairs
-                        go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2
-                        concatBms [] = VU.empty
-                        concatBms [(b, _v)] = b
-                        concatBms ((b1, v1) : (b2, v2) : rest') =
-                            let merged = go b1 (VB.length v1) b2 (VB.length v2)
-                             in concatBms ((merged, v1 <> v2) : rest')
-                     in Just $ concatBms (zip expandedBms allVecs)
+            newBm = concatValidity (zipWith Validity allBms (map VB.length allVecs))
          in BoxedColumn newBm (VB.concat allVecs)
     UnboxedColumn bm0 v0 ->
         let getCol (UnboxedColumn bm v) = case testEquality (typeOf v0) (typeOf v) of
@@ -891,18 +880,7 @@ concatManyColumns (c0 : cs) = case c0 of
             rest = map getCol cs
             allVecs = v0 : map snd rest
             allBms = bm0 : map fst rest
-            newBm
-                | all isNothing allBms = Nothing
-                | otherwise =
-                    let pairs = zip allVecs allBms
-                        expandedBms = map (\(v, mb) -> fromMaybe (allValidBitmap (VU.length v)) mb) pairs
-                        go b1 n1 b2 n2 = bitmapConcat n1 b1 n2 b2
-                        concatBms [] = VU.empty
-                        concatBms [(b, _)] = b
-                        concatBms ((b1, v1) : (b2, v2) : rest') =
-                            let merged = go b1 (VU.length v1) b2 (VU.length v2)
-                             in concatBms ((merged, v1 <> v2) : rest')
-                     in Just $ concatBms (zip expandedBms allVecs)
+            newBm = concatValidity (zipWith Validity allBms (map VU.length allVecs))
          in UnboxedColumn newBm (VU.concat allVecs)
     PackedText _ _ -> concatManyColumns (map materializePacked (c0 : cs))
     MergedColumn _ _ -> concatManyColumns (map materializeMerged (c0 : cs))

dataframe.cabal

@@ -316,6 +316,7 @@ test-suite tests
                    Functions,
                    GenDataFrame,
                    Internal.ColumnBuilder,
+                   Internal.ColumnConcat,
                    Internal.DictEncode,
                    Internal.Markdown,
                    Internal.PackedText,

tests/Internal/ColumnConcat.hs

@@ -0,0 +1,83 @@
+{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE TypeApplications #-}
+
+{- | Complexity regression for 'DI.concatManyColumns', the batch collector
+behind the lazy executor's @collectStream@. It exists to concatenate every
+chunk in one pass rather than folding @acc <> batch@, which would recopy the
+accumulator on every step.
+
+Allocation is deterministic where wall-clock is not, so its growth rate is a
+stable stand-in for the complexity.
+-}
+module Internal.ColumnConcat (tests) where
+
+import qualified Data.Vector as VB
+import qualified DataFrame.Internal.Column as DI
+
+import GHC.Conc.Sync (getAllocationCounter)
+import Test.HUnit
+
+-- | Nullable, so the concat has a bitmap to build as well as data to copy.
+chunk :: DI.Column
+chunk = DI.fromList (map Just [1 .. 250 :: Int])
+
+-- | Bytes allocated forcing a concat of @k@ chunks, bitmap included.
+allocationFor :: Int -> IO Integer
+allocationFor k = do
+    before <- getAllocationCounter
+    let !n = DI.numElements (DI.concatManyColumns (replicate k chunk))
+    after <- n `seq` getAllocationCounter
+    pure (fromIntegral (before - after))
+
+{- | The values and the null positions have to survive the concat, not just its
+allocation profile. Chunks are deliberately uneven and land off byte boundaries
+so a bitmap written at the wrong offset shows up.
+-}
+unevenChunks :: [DI.Column]
+unevenChunks =
+    [ nullableChunk 3 [1]
+    , nullableChunk 8 [0, 7]
+    , nullableChunk 5 []
+    , nullableChunk 1 [0]
+    , nullableChunk 11 [2, 10]
+    ]
+
+-- | @nullableChunk n nulls@ has values 0..n-1 with @nulls@ marked missing.
+nullableChunk :: Int -> [Int] -> DI.Column
+nullableChunk n nulls =
+    DI.fromVector
+        (VB.generate n (\i -> if i `elem` nulls then Nothing else Just i))
+
+concatManyColumnsMatchesChunks :: Test
+concatManyColumnsMatchesChunks =
+    TestCase
+        ( assertEqual
+            "concat preserves every value and null position"
+            (concat [DI.toList @(Maybe Int) c | c <- unevenChunks])
+            (DI.toList @(Maybe Int) (DI.concatManyColumns unevenChunks))
+        )
+
+{- | Quadrupling the chunk count quadruples the rows, so a single-pass concat
+allocates about 4x as much. Recopying the accumulator per chunk allocates about
+16x; the bound sits between the two.
+-}
+concatManyColumnsAllocatesLinearly :: Test
+concatManyColumnsAllocatesLinearly = TestCase $ do
+    small <- allocationFor 40
+    large <- allocationFor 160
+    assertBool
+        ( "allocation grew "
+            ++ show (large `div` max 1 small)
+            ++ "x for 4x the rows"
+        )
+        (large <= 8 * small)
+
+tests :: [Test]
+tests =
+    [ TestLabel
+        "concatManyColumns preserves values and nulls"
+        concatManyColumnsMatchesChunks
+    , TestLabel
+        "concatManyColumns allocates linearly"
+        concatManyColumnsAllocatesLinearly
+    ]

tests/Main.hs

@@ -14,6 +14,7 @@ import qualified IO.CsvGolden
 import qualified IO.JSON
 import qualified IR.ExprJsonRoundtrip
 import qualified Internal.ColumnBuilder
+import qualified Internal.ColumnConcat
 import qualified Internal.DictEncode
 import qualified Internal.Markdown
 import qualified Internal.PackedText
@@ -72,6 +73,7 @@ tests :: Test
 tests =
     TestList $
         Internal.ColumnBuilder.tests
+            ++ Internal.ColumnConcat.tests
             ++ Internal.DictEncode.tests
             ++ Internal.Markdown.tests
             ++ Internal.PackedText.tests