From bb139e11d0da47f954656aa2df5f0dfbdbaaa0dd Mon Sep 17 00:00:00 2001 From: Samuel Schlesinger Date: Thu, 20 Aug 2026 19:41:49 -0400 Subject: [PATCH 1/2] test: Pin concatManyColumns to its chunks and to linear allocation concatManyColumns concatenates the data in one VB.concat, but the nullable path folds the per-chunk bitmaps pairwise and carries the running length by concatenating the data vectors again at every step: concatBms ((b1, v1) : (b2, v2) : rest') = let merged = go b1 (VB.length v1) b2 (VB.length v2) in concatBms ((merged, v1 <> v2) : rest') That reinstates the accumulator recopy the function exists to avoid. Allocation grows 15x for 4x the rows where linear would be 4x. The second case is an oracle rather than a repro: uneven chunks that straddle byte boundaries must keep every value and null position, which is the property the new bitmap writes have to preserve. --- dataframe.cabal | 1 + tests/Internal/ColumnConcat.hs | 83 ++++++++++++++++++++++++++++++++++ tests/Main.hs | 2 + 3 files changed, 86 insertions(+) create mode 100644 tests/Internal/ColumnConcat.hs diff --git a/dataframe.cabal b/dataframe.cabal index 10f741dc..50c934a4 100644 --- a/dataframe.cabal +++ b/dataframe.cabal @@ -316,6 +316,7 @@ test-suite tests Functions, GenDataFrame, Internal.ColumnBuilder, + Internal.ColumnConcat, Internal.DictEncode, Internal.Markdown, Internal.PackedText, diff --git a/tests/Internal/ColumnConcat.hs b/tests/Internal/ColumnConcat.hs new file mode 100644 index 00000000..ecce36e1 --- /dev/null +++ b/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 + ] diff --git a/tests/Main.hs b/tests/Main.hs index 7689e2f7..020241b3 100644 --- a/tests/Main.hs +++ b/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 From 4a8209cb66a8f06f5b9a6da69164e39b2afe52a5 Mon Sep 17 00:00:00 2001 From: Samuel Schlesinger Date: Thu, 20 Aug 2026 19:41:49 -0400 Subject: [PATCH 2/2] fix: Concatenate chunk bitmaps in a single pass The nullable path folded the per-chunk bitmaps pairwise and concatenated the data vectors on every step just to carry a running length, so it recopied the accumulator exactly the way the function's one-pass design sets out to avoid. Take the lengths up front and splice every chunk's bits into one preallocated destination with concatValidity, the byte-level splice the Bitmap module already has. For 160 chunks of 10k nullable rows: allocation growth drops from 15x to 4x per 4x the rows, and forcing the bitmap from 449ms to 7.9ms. --- .../DataFrame/Internal/Column/Operations.hs | 28 ++----------------- 1 file changed, 3 insertions(+), 25 deletions(-) diff --git a/dataframe-core/src-internal/DataFrame/Internal/Column/Operations.hs b/dataframe-core/src-internal/DataFrame/Internal/Column/Operations.hs index 750ac9a7..c82cfc85 100644 --- a/dataframe-core/src-internal/DataFrame/Internal/Column/Operations.hs +++ b/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))