Merge pull request #489 from IntersectMBO/jdral/n-way-unions

dcoutts · web-flow · commit e256e9b8bf4b · 2024-12-10T16:29:54.000Z
Add n-way table unions
diff --git a/src/Database/LSMTree.hs b/src/Database/LSMTree.hs
@@ -80,6 +80,7 @@ module Database.LSMTree (
 
     -- * Table union
   , union
+  , unions
 
     -- * Serialisation
   , SerialiseKey
@@ -530,6 +531,16 @@ union :: forall m k v b.
   -> m (Table m k v b)
 union = error "union: not yet implemented" $ union @m @k @v @b
 
+{-# SPECIALISE unions ::
+     ResolveValue v
+  => V.Vector (Table IO k v b)
+  -> IO (Table IO k v b) #-}
+unions :: forall m k v b.
+     (IOLike m, ResolveValue v)
+  => V.Vector (Table m k v b)
+  -> m (Table m k v b)
+unions = error "unions: not yet implemented" $ unions @m @k @v
+
 {-------------------------------------------------------------------------------
   Monoidal value resolution
 -------------------------------------------------------------------------------}
diff --git a/src/Database/LSMTree/Monoidal.hs b/src/Database/LSMTree/Monoidal.hs
@@ -97,6 +97,7 @@ module Database.LSMTree.Monoidal (
 
     -- * Table union
   , union
+  , unions
 
     -- * Concurrency
     -- $concurrency
@@ -672,6 +673,24 @@ union :: forall m k v.
   -> m (Table m k v)
 union = error "union: not yet implemented" $ union @m @k @v
 
+{-# SPECIALISE unions ::
+     ResolveValue v
+  => V.Vector (Table IO k v)
+  -> IO (Table IO k v) #-}
+-- | Like 'union', but for @n@ tables.
+--
+-- A good mental model of this operation is @'Data.Map.Lazy.unionsWith' (<>)@ on
+-- @'Data.Map.Lazy.Map' k v@.
+--
+-- Exceptions:
+--
+-- * Unioning 0 tables is an exception.
+unions :: forall m k v.
+     (IOLike m, ResolveValue v)
+  => V.Vector (Table m k v)
+  -> m (Table m k v)
+unions = error "unions: not yet implemented" $ unions @m @k @v
+
 {-------------------------------------------------------------------------------
   Monoidal value resolution
 -------------------------------------------------------------------------------}
diff --git a/src/Database/LSMTree/Normal.hs b/src/Database/LSMTree/Normal.hs
@@ -98,6 +98,7 @@ module Database.LSMTree.Normal (
 
     -- * Table union
   , union
+  , unions
 
     -- * Concurrency #concurrency#
     -- $concurrency
@@ -787,3 +788,18 @@ union :: forall m k v b.
   -> Table m k v b
   -> m (Table m k v b)
 union = error "union: not yet implemented" $ union @m @k @v
+
+{-# SPECIALISE unions :: V.Vector (Table IO k v b) -> IO (Table IO k v b) #-}
+-- | Like 'union', but for @n@ tables.
+--
+-- A good mental model of this operation is @'Data.Map.Lazy.unions'@ on
+-- @'Data.Map.Lazy.Map' k v@.
+--
+-- Exceptions:
+--
+-- * Unioning 0 tables is an exception.
+unions :: forall m k v b.
+     IOLike m
+  => V.Vector (Table m k v b)
+  -> m (Table m k v b)
+unions = error "union: not yet implemented" $ union @m @k @v
diff --git a/test/Database/LSMTree/Class.hs b/test/Database/LSMTree/Class.hs
@@ -8,6 +8,7 @@ module Database.LSMTree.Class (
   , withTableFromSnapshot
   , withTableDuplicate
   , withTableUnion
+  , withTableUnions
   , withCursor
   , module Common
   , module Types
@@ -162,6 +163,13 @@ class (IsSession (Session h)) => IsTable h where
         -> h m k v b
         -> m (h m k v b)
 
+    unions ::
+           ( IOLike m
+           , C k v b
+           )
+        => V.Vector (h m k v b)
+        -> m (h m k v b)
+
 withTableNew :: forall h m k v b a.
     (IOLike m, IsTable h, C k v b)
   => Session h m
@@ -194,6 +202,13 @@ withTableUnion :: forall h m k v b a.
   -> m a
 withTableUnion table1 table2 = bracket (table1 `union` table2) close
 
+withTableUnions :: forall h m k v b a.
+     (IOLike m, IsTable h, C k v b)
+  => V.Vector (h m k v b)
+  -> (h m k v b -> m a)
+  -> m a
+withTableUnions tables = bracket (unions tables) close
+
 withCursor :: forall h m k v b a.
      (IOLike m, IsTable h, C k v b)
   => Maybe k
@@ -232,3 +247,4 @@ instance IsTable R.Table where
 
     duplicate = R.duplicate
     union = R.union
+    unions = R.unions
diff --git a/test/Database/LSMTree/Model/IO.hs b/test/Database/LSMTree/Model/IO.hs
@@ -17,25 +17,26 @@ module Database.LSMTree.Model.IO (
 import           Control.Concurrent.Class.MonadSTM.Strict
 import           Control.Exception (Exception)
 import           Control.Monad.Class.MonadThrow (MonadThrow (..))
+import qualified Data.Vector as V
 import qualified Database.LSMTree.Class as Class
 import           Database.LSMTree.Model.Session (TableConfig (..))
 import qualified Database.LSMTree.Model.Session as Model
 
 newtype Session m = Session (StrictTVar m (Maybe Model.Model))
 
 data Table m k v b = Table {
-    _thSession :: !(Session m)
-  , _thTable   :: !(Model.Table k v b)
+    thSession :: !(Session m)
+  , thTable   :: !(Model.Table k v b)
   }
 
 data BlobRef m b = BlobRef {
-    _brSession :: !(Session m)
-  , _brBlobRef :: !(Model.BlobRef b)
+    brSession :: !(Session m)
+  , brBlobRef :: !(Model.BlobRef b)
   }
 
 data Cursor m k v b = Cursor {
-    _cSession :: !(Session m)
-  , _cCursor  :: !(Model.Cursor k v b)
+    cSession :: !(Session m)
+  , cCursor  :: !(Model.Cursor k v b)
   }
 
 newtype Err = Err (Model.Err)
@@ -76,7 +77,7 @@ instance Class.IsTable Table where
 
     rangeLookup (Table s t) x1 = fmap (fmap (BlobRef s)) <$>
       runInOpenSession s (Model.rangeLookup x1 t)
-    retrieveBlobs _ s x1 = runInOpenSession s (Model.retrieveBlobs (fmap _brBlobRef x1))
+    retrieveBlobs _ s x1 = runInOpenSession s (Model.retrieveBlobs (fmap brBlobRef x1))
 
     newCursor k (Table s t) = Cursor s <$> runInOpenSession s (Model.newCursor k t)
     closeCursor _ (Cursor s c) = runInOpenSession s (Model.closeCursor c)
@@ -90,3 +91,8 @@ instance Class.IsTable Table where
 
     union (Table s1 t1) (Table _s2 t2) =
       Table s1 <$> runInOpenSession s1 (Model.union Model.getResolve t1 t2)
+
+    unions ts =
+        Table s <$> runInOpenSession s (Model.unions Model.getResolve (V.map thTable ts))
+      where
+        Table s _ = V.head ts
diff --git a/test/Database/LSMTree/Model/Session.hs b/test/Database/LSMTree/Model/Session.hs
@@ -73,6 +73,7 @@ module Database.LSMTree.Model.Session (
   , duplicate
     -- * Table union
   , union
+  , unions
   ) where
 
 import           Control.Monad (when)
@@ -204,6 +205,8 @@ data Err =
   | ErrSnapshotWrongType
   | ErrBlobRefInvalidated
   | ErrCursorClosed
+    -- | Passed zero tables to 'unions'
+  | ErrUnionsZeroTables
   deriving stock (Show, Eq)
 
 {-------------------------------------------------------------------------------
@@ -628,3 +631,21 @@ union r th1 th2 = do
   (_, t1) <- guardTableIsOpen th1
   (_, t2) <- guardTableIsOpen th2
   newTableWith TableConfig $ Model.union r t1 t2
+
+unions ::
+     ( MonadState Model m
+     , MonadError Err m
+     , C k v b
+     )
+  => ResolveSerialisedValue v
+  -> V.Vector (Table k v b)
+  -> m (Table k v b)
+unions r tables
+  | n == 0 = throwError ErrUnionsZeroTables
+  | otherwise = do
+      tables' <- V.forM tables $ \table -> do
+        (_, table') <- guardTableIsOpen table
+        pure table'
+      newTableWith TableConfig $ Model.unions r tables'
+  where
+    n = V.length tables
diff --git a/test/Database/LSMTree/Model/Table.hs b/test/Database/LSMTree/Model/Table.hs
@@ -38,6 +38,7 @@ module Database.LSMTree.Model.Table (
   , duplicate
     -- * Table union
   , union
+  , unions
     -- * Testing
   , size
   ) where
@@ -301,24 +302,22 @@ readCursor n c =
     )
 
 {-------------------------------------------------------------------------------
-  Merging tables
+  Table union
 -------------------------------------------------------------------------------}
 
--- | Merge full tables, creating a new table.
---
--- NOTE: close tables using 'close' as soon as they are
--- unused.
---
--- Multiple tables of the same type but with different configuration parameters
--- can live in the same session. However, some operations, like
+-- | Union two full tables, creating a new table.
 union ::
      ResolveSerialisedValue v
   -> Table k v b
   -> Table k v b
   -> Table k v b
 union r (Table xs) (Table ys) =
-    Table (Map.unionWith f xs ys)
-  where
-    f (v1, bMay1) (v2, bMay2) =
-      (resolveSerialised r v1 v2, getFirst (First bMay1 <> First bMay2))
+    Table (Map.unionWith (resolveValueAndBlob r) xs ys)
 
+-- | Like 'union', but for @n@ tables.
+unions ::
+     ResolveSerialisedValue v
+  -> V.Vector (Table k v b)
+  -> Table k v b
+unions r tables =
+    Table (Map.unionsWith (resolveValueAndBlob r) (V.map values tables))
diff --git a/test/Test/Database/LSMTree/StateMachine.hs b/test/Test/Database/LSMTree/StateMachine.hs
diff --git a/test/Test/Database/LSMTree/UnitTests.hs b/test/Test/Database/LSMTree/UnitTests.hs
