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)) 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