Merge pull request #536 from IntersectMBO/mheinzel/merge-union

mheinzel · web-flow · commit 3e6bd8da2658 · 2025-01-22T17:59:07.000Z
Implement MergeUnion merges
diff --git a/bench/micro/Bench/Database/LSMTree/Internal/Merge.hs b/bench/micro/Bench/Database/LSMTree/Internal/Merge.hs
@@ -84,23 +84,39 @@ benchmarks = bgroup "Bench.Database.LSMTree.Internal.Merge" [
                          randomWord64OutOf (totalEntries `div` 2)
         , mergeMappend = Just (onDeserialisedValues ((+) @Word64))
         }
+      -- different merge types
     , benchMerge configWord64
-        { name         = "word64-mix-x4"
+        { name         = "word64-mix-collisions-x4-midlevel"
         , nentries     = totalEntries `splitInto` 4
         , finserts     = 1
         , fdeletes     = 1
         , fmupserts    = 1
+        , randomKey    = -- each run uses half of the possible keys
+                         randomWord64OutOf (totalEntries `div` 2)
+        , mergeMappend = Just (onDeserialisedValues ((+) @Word64))
+        , mergeType    = MergeMidLevel
+        }
+    , benchMerge configWord64
+        { name         = "word64-mix-collisions-x4-lastlevel"
+        , nentries     = totalEntries `splitInto` 4
+        , finserts     = 1
+        , fdeletes     = 1
+        , fmupserts    = 1
+        , randomKey    = -- each run uses half of the possible keys
+                         randomWord64OutOf (totalEntries `div` 2)
         , mergeMappend = Just (onDeserialisedValues ((+) @Word64))
+        , mergeType    = MergeLastLevel
         }
     , benchMerge configWord64
-        { name         = "word64-mix-collisions-x4"
+        { name         = "word64-mix-collisions-x4-union"
         , nentries     = totalEntries `splitInto` 4
         , finserts     = 1
         , fdeletes     = 1
         , fmupserts    = 1
         , randomKey    = -- each run uses half of the possible keys
                          randomWord64OutOf (totalEntries `div` 2)
         , mergeMappend = Just (onDeserialisedValues ((+) @Word64))
+        , mergeType    = MergeUnion
         }
       -- not writing anything at all
     , benchMerge configWord64
@@ -174,6 +190,20 @@ benchmarks = bgroup "Bench.Database.LSMTree.Internal.Merge" [
         , fdeletes     = 1
         , mergeType    = MergeLastLevel
         }
+    , benchMerge configUTxOStaking
+        { name         = "utxo-x2-tree-union"  -- binary union merge
+        , nentries     = totalEntries `distributed` [4, 1]
+        , mergeType    = MergeUnion
+        }
+    , benchMerge configUTxOStaking
+        { name         = "utxo-x10-tree-level"  -- merge a whole table (for union)
+        , nentries     = totalEntries `distributed` [ 1, 1, 1
+                                                    , 4, 4, 4
+                                                    , 16, 16, 16
+                                                    , 100  -- last level
+                                                    ]
+        , mergeType    = MergeLastLevel
+        }
     ]
   where
     totalEntries = 50_000
@@ -311,6 +341,14 @@ configUTxO = defaultConfig {
   , randomValue  = first serialiseValue . uniform @_ @UTxOValue
   }
 
+configUTxOStaking :: Config
+configUTxOStaking = defaultConfig {
+    fmupserts    = 1
+  , randomKey    = first serialiseKey . uniform @_ @UTxOKey
+  , randomValue  = first serialiseValue . uniform @_ @Word64
+  , mergeMappend = Just (onDeserialisedValues ((+) @Word64))
+  }
+
 mergeEnv ::
      Config
   -> IO ( FilePath -- ^ Temporary directory
diff --git a/src-extras/Database/LSMTree/Extras/Generators.hs b/src-extras/Database/LSMTree/Extras/Generators.hs
@@ -556,8 +556,7 @@ instance Arbitrary BlobSpan where
 -------------------------------------------------------------------------------}
 
 instance Arbitrary MergeType where
-  arbitrary = QC.elements [MergeMidLevel, MergeLastLevel]
-                          -- TODO: add MergeUnion once it is supported.
+  arbitrary = QC.elements [MergeMidLevel, MergeLastLevel, MergeUnion]
   shrink MergeMidLevel  = []
   shrink MergeLastLevel = [MergeMidLevel]
   shrink MergeUnion     = [MergeLastLevel]
diff --git a/src/Database/LSMTree/Internal/Entry.hs b/src/Database/LSMTree/Internal/Entry.hs
@@ -7,6 +7,7 @@ module Database.LSMTree.Internal.Entry (
   , unNumEntries
     -- * Value resolution/merging
   , combine
+  , combineUnion
   , combineMaybe
   ) where
 
@@ -105,6 +106,24 @@ combine f   (Mupdate u)       (Insert v)              = Insert (f u v)
 combine f   (Mupdate u)       (InsertWithBlob v blob) = InsertWithBlob (f u v) blob
 combine f   (Mupdate u)       (Mupdate v)             = Mupdate (f u v)
 
+-- | Combine two entries of runs that have been 'union'ed together. If any one
+-- has a value, the result should have a value (represented by 'Insert'). If
+-- both have a value, these values get combined monoidally.
+combineUnion :: (v -> v -> v) -> Entry v b -> Entry v b -> Entry v b
+combineUnion f = go
+  where
+    go Delete               e                    = e
+    go e                    Delete               = e
+    go (Insert u)           (Insert v)           = Insert (f u v)
+    go (Insert u)           (InsertWithBlob v b) = InsertWithBlob (f u v) b
+    go (Insert u)           (Mupdate v)          = Insert (f u v)
+    go (InsertWithBlob u b) (Insert v)           = InsertWithBlob (f u v) b
+    go (InsertWithBlob u b) (InsertWithBlob v _) = InsertWithBlob (f u v) b
+    go (InsertWithBlob u b) (Mupdate v)          = InsertWithBlob (f u v) b
+    go (Mupdate u)          (Insert v)           = Insert (f u v)
+    go (Mupdate u)          (InsertWithBlob v b) = InsertWithBlob (f u v) b
+    go (Mupdate u)          (Mupdate v)          = Insert (f u v)
+
 combineMaybe :: (v -> v -> v) -> Maybe (Entry v b) -> Maybe (Entry v b) -> Maybe (Entry v b)
 combineMaybe _ e1 Nothing          = e1
 combineMaybe _ Nothing e2          = e2
diff --git a/src/Database/LSMTree/Internal/Merge.hs b/src/Database/LSMTree/Internal/Merge.hs
@@ -252,22 +252,35 @@ steps ::
   => Merge m h
   -> Int  -- ^ How many input entries to consume (at least)
   -> m (Int, StepResult)
-steps Merge {..} requestedSteps = assertStepsInvariant <$> do
+steps m@Merge {..} requestedSteps = assertStepsInvariant <$> do
     -- TODO: ideally, we would not check whether the merge was already done on
     -- every call to @steps@. It is important for correctness, however, that we
     -- do not call @steps@ on a merge when it was already done. It is not yet
     -- clear whether our (upcoming) implementation of scheduled merges is going
     -- to satisfy this precondition when it calls @steps@, so for now we do the
     -- check.
     readMutVar mergeState >>= \case
-      Merging -> go 0
+      Merging -> case mergeType of
+        MergeMidLevel  -> doStepsLevel m requestedSteps
+        MergeLastLevel -> doStepsLevel m requestedSteps
+        MergeUnion     -> doStepsUnion m requestedSteps
       MergingDone -> pure (0, MergeDone)
       Completed -> error "steps: Merge is completed"
       Closed -> error "steps: Merge is closed"
   where
     assertStepsInvariant res = assert (stepsInvariant requestedSteps res) res
 
-    go :: Int -> m (Int, StepResult)
+{-# SPECIALISE doStepsLevel ::
+     Merge IO h
+  -> Int
+  -> IO (Int, StepResult) #-}
+doStepsLevel ::
+     (MonadMask m, MonadSTM m, MonadST m)
+  => Merge m h
+  -> Int  -- ^ How many input entries to consume (at least)
+  -> m (Int, StepResult)
+doStepsLevel Merge {..} requestedSteps = go 0
+  where
     go !n
       | n >= requestedSteps =
           return (n, MergeInProgress)
@@ -329,6 +342,48 @@ steps Merge {..} requestedSteps = assertStepsInvariant <$> do
             writeMutVar mergeState $! MergingDone
             pure (n + dropped, MergeDone)
 
+{-# SPECIALISE doStepsUnion ::
+     Merge IO h
+  -> Int
+  -> IO (Int, StepResult) #-}
+doStepsUnion ::
+     (MonadMask m, MonadSTM m, MonadST m)
+  => Merge m h
+  -> Int  -- ^ How many input entries to consume (at least)
+  -> m (Int, StepResult)
+doStepsUnion Merge {..} requestedSteps = go 0
+  where
+    go !n
+      | n >= requestedSteps =
+          return (n, MergeInProgress)
+      | otherwise = do
+          (key, entry, hasMore) <- Readers.pop mergeReaders
+          handleEntry (n + 1) key entry hasMore
+
+    -- Similar to 'handleMupdate' in 'stepsLevel', but here we have to combine
+    -- all entries monoidally, so there are no obsolete/overwritten entries
+    -- that we could skip.
+    handleEntry !n !key !entry Readers.Drained = do
+        -- no future entries, no previous entry to resolve, just write!
+        writeReaderEntry mergeType mergeBuilder key entry
+        writeMutVar mergeState $! MergingDone
+        pure (n, MergeDone)
+
+    handleEntry !n !key !entry Readers.HasMore = do
+        nextKey <- Readers.peekKey mergeReaders
+        if nextKey /= key
+          then do
+            -- resolved all entries for this key, write it
+            writeReaderEntry mergeType mergeBuilder key entry
+            go n
+          else do
+            (_, nextEntry, hasMore) <- Readers.pop mergeReaders
+            -- for resolution, we need the full second value to be present
+            let resolved = combineUnion mergeMappend
+                             (Reader.toFullEntry entry)
+                             (Reader.toFullEntry nextEntry)
+            handleEntry (n + 1) key (Reader.Entry resolved) hasMore
+
 {-# SPECIALISE writeReaderEntry ::
      MergeType
   -> RunBuilder IO h
diff --git a/test/Test/Database/LSMTree/Internal/Merge.hs b/test/Test/Database/LSMTree/Internal/Merge.hs
@@ -1,5 +1,3 @@
-{-# LANGUAGE LambdaCase #-}
-
 module Test.Database.LSMTree.Internal.Merge (tests) where
 
 import           Control.Exception (evaluate)
@@ -9,7 +7,7 @@ import qualified Data.BloomFilter as Bloom
 import           Data.Foldable (traverse_)
 import           Data.Map.Strict (Map)
 import qualified Data.Map.Strict as Map
-import           Data.Maybe (isJust)
+import           Data.Maybe (isJust, mapMaybe)
 import qualified Data.Vector as V
 import           Database.LSMTree.Extras
 import           Database.LSMTree.Extras.Generators (KeyForIndexCompact)
@@ -38,12 +36,15 @@ import           Test.Tasty.QuickCheck
 
 tests :: TestTree
 tests = testGroup "Test.Database.LSMTree.Internal.Merge"
-    [ testProperty "prop_MergeDistributes" $ \level stepSize wbs ->
+    [ testProperty "prop_MergeDistributes" $ \mergeType stepSize rds ->
+        ioPropertyWithMockFS $ \fs hbio ->
+          prop_MergeDistributes fs hbio mergeType stepSize rds
+    , testProperty "prop_MergeUnion" $ \stepSize rds ->
         ioPropertyWithMockFS $ \fs hbio ->
-          prop_MergeDistributes fs hbio level stepSize wbs
-    , testProperty "prop_AbortMerge" $ \level stepSize wbs ->
+          prop_MergeUnion fs hbio stepSize rds
+    , testProperty "prop_AbortMerge" $ \level stepSize rds ->
         ioPropertyWithMockFS $ \fs hbio ->
-          prop_AbortMerge fs hbio level stepSize wbs
+          prop_AbortMerge fs hbio level stepSize rds
     ]
   where
     ioPropertyWithMockFS ::
@@ -70,12 +71,12 @@ prop_MergeDistributes ::
      StepSize ->
      SmallList (RunData KeyForIndexCompact SerialisedValue SerialisedBlob) ->
      IO Property
-prop_MergeDistributes fs hbio level stepSize (SmallList rds) =
+prop_MergeDistributes fs hbio mergeType stepSize (SmallList rds) =
     withRuns fs hbio (V.fromList (zip (simplePaths [10..]) rds')) $ \runs -> do
       let stepsNeeded = sum (map (Map.size . unRunData) rds)
-      (stepsDone, lhs) <- mergeRuns fs hbio level (RunNumber 0) runs stepSize
-      withRun fs hbio (simplePath 1)
-              (RunData $ mergeWriteBuffers level $ fmap unRunData rds') $ \rhs -> do
+      (stepsDone, lhs) <- mergeRuns fs hbio mergeType (RunNumber 0) runs stepSize
+      let runData = RunData $ mergeWriteBuffers mergeType $ fmap unRunData rds'
+      withRun fs hbio (simplePath 1) runData $ \rhs -> do
 
         (lhsSize, lhsFilter, lhsIndex, lhsKOps,
          lhsKOpsFileContent, lhsBlobFileContent) <- getRunContent lhs
@@ -133,6 +134,45 @@ prop_MergeDistributes fs hbio level stepSize (SmallList rds) =
              , blobFileContent
              )
 
+-- | Union-merging multiple runs behaves like 'Map.unionsWith' on their values
+-- and blobs.
+prop_MergeUnion ::
+     FS.HasFS IO h ->
+     FS.HasBlockIO IO h ->
+     StepSize ->
+     SmallList (RunData KeyForIndexCompact SerialisedValue SerialisedBlob) ->
+     IO Property
+prop_MergeUnion fs hbio stepSize (SmallList rds) =
+    withRuns fs hbio (V.fromList (zip (simplePaths [10..]) rds')) $ \runs -> do
+      (_, run) <- mergeRuns fs hbio MergeUnion (RunNumber 0) runs stepSize
+
+      lhsKOps <- readKOps Nothing run
+      let lhs = Map.fromList (mapMaybe (traverse getValueAndBlob) lhsKOps)
+
+      -- cleanup
+      releaseRef run
+
+      return $
+             lhs === rhs
+        .&&. counterexample ("Deletes in " <> show lhs)
+               (all ((/= Entry.Delete) . snd) lhsKOps)
+  where
+    rds' = fmap serialiseRunData rds
+
+    rhs :: Map SerialisedKey (SerialisedValue, Maybe SerialisedBlob)
+    rhs = Map.unionsWith resolveValueAndBlob
+            (map (Map.mapMaybe getValueAndBlob . unRunData) rds')
+
+    getValueAndBlob :: Entry.Entry v b -> Maybe (v, Maybe b)
+    getValueAndBlob = \case
+        Entry.Insert v           -> Just (v, Nothing)
+        Entry.InsertWithBlob v b -> Just (v, Just b)
+        Entry.Mupdate v          -> Just (v, Nothing)
+        Entry.Delete             -> Nothing
+
+    resolveValueAndBlob (v', Nothing) (v, b) = (mappendValues v' v, b)
+    resolveValueAndBlob (v', Just b)  (v, _) = (mappendValues v' v, Just b)
+
 -- | After merging for a few steps, we can prematurely abort the merge, which
 -- should clean up properly.
 prop_AbortMerge ::
@@ -196,10 +236,12 @@ type SerialisedEntry = Entry.Entry SerialisedValue SerialisedBlob
 mergeWriteBuffers :: MergeType
                   -> [Map SerialisedKey SerialisedEntry]
                   ->  Map SerialisedKey SerialisedEntry
-mergeWriteBuffers mergeType =
---TODO: review and update this to support MergeUnion
-        (if mergeType == MergeMidLevel then id else Map.filter (not . isDelete))
-      . Map.unionsWith (Entry.combine mappendValues)
+mergeWriteBuffers = \case
+    MergeMidLevel  -> Map.unionsWith (Entry.combine mappendValues)
+    MergeLastLevel -> Map.filter (not . isDelete)
+                    . Map.unionsWith (Entry.combine mappendValues)
+    MergeUnion     -> Map.filter (not . isDelete)
+                    . Map.unionsWith (Entry.combineUnion mappendValues)
   where
     isDelete Entry.Delete = True
     isDelete _            = False
