Improve querying by skipping blocks (#173)

* Optimize queryEntry to skip records

* Update test spec

* Rename decodeChunkEntries to streamChunkEntries

* Rename recursive function variable

Author: joaosreis

scls-format/src/Cardano/SCLS/Internal/Reader.hs

@@ -19,7 +19,7 @@ module Cardano.SCLS.Internal.Reader (
   knownNamespacedData,
 
   -- * Low-level functions
-  decodeChunkEntries,
+  streamChunkEntries,
 ) where
 
 import Cardano.SCLS.Internal.Frame
@@ -51,13 +51,11 @@ import GHC.TypeLits (KnownSymbol)
 import Streaming qualified as S
 import Streaming.Prelude qualified as S
 
-{- | Decode entries from chunk data into a stream.
-
-This function provides a stream of the Chunk entries
+{- | This function provides a stream of the Chunk entries
 stored in the data field of the 'Chunk'.
 -}
-decodeChunkEntries :: (Typeable u, MemPack u) => BS.ByteString -> S.Stream (S.Of u) IO ()
-decodeChunkEntries = go
+streamChunkEntries :: (Typeable u, MemPack u) => BS.ByteString -> S.Stream (S.Of u) IO ()
+streamChunkEntries = go
  where
   go !bs
     | BS.null bs = pure ()
@@ -170,7 +168,7 @@ namespacedData handle namespace = stream
           fetchOffsetFrame handle next_record
         for_ (decodeFrame dataRecord) \chunkRecord -> do
           when (chunkNamespace (frameViewContent (chunkRecord)) == namespace) do
-            decodeChunkEntries (chunkData $ frameViewContent chunkRecord)
+            streamChunkEntries (chunkData $ frameViewContent chunkRecord)
         go next_record
 
 knownNamespacedData :: forall ns. (KnownSymbol ns, KnownNamespace ns) => Handle -> Proxy ns -> S.Stream (S.Of (ChunkEntry (NamespaceKey ns) (NamespaceEntry ns))) IO ()

scls-format/src/Cardano/SCLS/Query.hs

@@ -1,42 +1,68 @@
+{-# LANGUAGE BlockArguments #-}
 {-# LANGUAGE LambdaCase #-}
 
 module Cardano.SCLS.Query (queryEntry) where
 
-import Cardano.SCLS.Internal.Reader (withNamespacedData)
+import Cardano.SCLS.Internal.Frame (FrameView (frameViewContent), decodeFrame, fetchNextFrame, fetchOffsetFrame, headerOffset)
+import Cardano.SCLS.Internal.Reader (streamChunkEntries)
+import Cardano.SCLS.Internal.Record.Chunk (Chunk (chunkData, chunkNamespace))
 import Cardano.SCLS.Internal.Serializer.HasKey (HasKey (..))
 import Cardano.Types.Namespace (Namespace)
+import Control.Monad (mfilter)
+import Control.Monad.Fix (fix)
+import Data.Functor ((<&>))
 import Data.MemPack (MemPack)
 import Data.Typeable (Typeable)
 import Streaming.Prelude qualified as S
+import System.IO (IOMode (ReadMode), SeekMode (AbsoluteSeek), hSeek, withBinaryFile)
 
 {- | Query for a specific entry by namespace and key.
 
 This function searches for an entry with the given key within the specified
-namespace in the SCLS data directory. It leverages the assumption that entries
-are stored in sorted order by key, allowing for early termination once a key
-greater than or equal to the target is found.
+namespace and SCLS file. It leverages the assumption that entries are stored
+in sorted order by key, allowing for early termination once a key greater than
+or equal to the target is found.
 -}
 queryEntry ::
   (HasKey a, MemPack a, Typeable a) =>
-  -- | The directory path containing the SCLS data files
+  -- | The path for the SCLS file
   FilePath ->
   -- | The namespace to search within
   Namespace ->
   -- | The key of the entry to find
   Key a ->
   -- | Returns @Just entry@ if an entry with the exact key is found, @Nothing@ otherwise
   IO (Maybe a)
-queryEntry dir ns key = do
-  withNamespacedData dir ns $ \stream -> do
-    value :: Maybe a <-
-      S.head_ $
-        S.dropWhile
-          ( \entry ->
-              let k = getKey entry
-               in k < key
-          )
-          stream
-    case value of
-      Just entry | getKey entry == key -> pure $ Just entry
-      Nothing -> pure Nothing
-      Just _ -> pure Nothing
+queryEntry filePath namespace key = do
+  withBinaryFile
+    filePath
+    ReadMode
+    query
+ where
+  query handle = do
+    hSeek handle AbsoluteSeek 0
+    flip fix (headerOffset, S.each []) \go (offset, prev) -> do
+      fetchNextFrame handle offset >>= \case
+        Nothing ->
+          -- No more records. If the entry exists, it must be on the previous record.
+          findInStream prev
+        Just nextOffset -> do
+          dataRecord <- fetchOffsetFrame handle nextOffset
+          case frameViewContent <$> decodeFrame dataRecord of
+            Right chunk
+              | chunkNamespace chunk == namespace -> do
+                  let entries = streamChunkEntries (chunkData chunk)
+                  S.head_ entries >>= \case
+                    Just entry
+                      | (getKey entry == key) -> pure $ Just entry
+                      | (getKey entry > key) ->
+                          -- The key of the current block's first entry is bigger than the queried key.
+                          -- If the entry exists, it must be on the previous record.
+                          findInStream prev
+                      | otherwise -> go (nextOffset, entries)
+                    Nothing -> go (nextOffset, prev)
+            _ -> go (nextOffset, prev)
+
+  findInStream stream = do
+    S.head_ (S.dropWhile ((< key) . getKey) stream)
+      <&> mfilter ((== key) . getKey)

scls-format/test/QuerySpec.hs

@@ -80,8 +80,7 @@ basicTests =
 
         serializeTestData fileName [(ns, entries)]
 
-        result <- queryEntry fileName ns key
-        result `shouldBe` Just (mkTestEntry 100 1 value)
+        queryEntry fileName ns key `shouldReturn` Just (mkTestEntry 100 1 value)
 
     it "returns Nothing when entry doesn't exist" $
       withSystemTempDirectory "query-test-XXXXXX" $ \tmpDir -> do
@@ -92,8 +91,7 @@ basicTests =
 
         serializeTestData fileName [(ns, entries)]
 
-        result <- queryEntry @TestEntry fileName ns searchKey
-        result `shouldBe` Nothing
+        queryEntry @TestEntry fileName ns searchKey `shouldReturn` Nothing
 
     it "finds entry among multiple entries" $
       withSystemTempDirectory "query-test-XXXXXX" $ \tmpDir -> do
@@ -109,8 +107,7 @@ basicTests =
 
         serializeTestData fileName [(ns, entries)]
 
-        result <- queryEntry fileName ns searchKey
-        result `shouldBe` Just (mkTestEntry 150 5 value)
+        queryEntry fileName ns searchKey `shouldReturn` Just (mkTestEntry 150 5 value)
 
 edgeCaseTests :: Spec
 edgeCaseTests =
@@ -123,8 +120,7 @@ edgeCaseTests =
 
         serializeTestData fileName [(ns, [])]
 
-        result <- queryEntry @TestEntry fileName ns searchKey
-        result `shouldBe` Nothing
+        queryEntry @TestEntry fileName ns searchKey `shouldReturn` Nothing
 
     it "returns Nothing for non-existent namespace" $
       withSystemTempDirectory "query-test-XXXXXX" $ \tmpDir -> do
@@ -136,8 +132,7 @@ edgeCaseTests =
 
         serializeTestData fileName [(existingNs, entries)]
 
-        result <- queryEntry @TestEntry fileName nonExistentNs searchKey
-        result `shouldBe` Nothing
+        queryEntry @TestEntry fileName nonExistentNs searchKey `shouldReturn` Nothing
 
     it "returns first match when there are duplicate keys" $
       withSystemTempDirectory "query-test-XXXXXX" $ \tmpDir -> do
@@ -153,8 +148,7 @@ edgeCaseTests =
 
         serializeTestData fileName [(ns, entries)]
 
-        result <- queryEntry fileName ns searchKey
-        result `shouldBe` Just (mkTestEntry 100 1 firstValue)
+        queryEntry fileName ns searchKey `shouldReturn` Just (mkTestEntry 100 1 firstValue)
 
     it "finds single entry in namespace" $
       withSystemTempDirectory "query-test-XXXXXX" $ \tmpDir -> do
@@ -166,8 +160,7 @@ edgeCaseTests =
 
         serializeTestData fileName [(ns, entries)]
 
-        result <- queryEntry fileName ns searchKey
-        result `shouldBe` Just (mkTestEntry 42 7 value)
+        queryEntry fileName ns searchKey `shouldReturn` Just (mkTestEntry 42 7 value)
 
 multiNamespaceTests :: Spec
 multiNamespaceTests =
@@ -185,11 +178,9 @@ multiNamespaceTests =
 
         serializeTestData fileName [(ns1, entries1), (ns2, entries2)]
 
-        result1 <- queryEntry fileName ns1 searchKey
-        result1 `shouldBe` Just (mkTestEntry 100 1 value1)
+        queryEntry fileName ns1 searchKey `shouldReturn` Just (mkTestEntry 100 1 value1)
 
-        result2 <- queryEntry fileName ns2 searchKey
-        result2 `shouldBe` Just (mkTestEntry 100 1 value2)
+        queryEntry fileName ns2 searchKey `shouldReturn` Just (mkTestEntry 100 1 value2)
 
     it "correctly isolates namespaces - entry exists in one but not another" $
       withSystemTempDirectory "query-test-XXXXXX" $ \tmpDir -> do
@@ -202,11 +193,9 @@ multiNamespaceTests =
 
         serializeTestData fileName [(ns1, entries1), (ns2, entries2)]
 
-        result1 <- queryEntry fileName ns1 searchKey
-        result1 `shouldBe` Just (mkTestEntry 100 1 "exists-here")
+        queryEntry fileName ns1 searchKey `shouldReturn` Just (mkTestEntry 100 1 "exists-here")
 
-        result2 <- queryEntry @TestEntry fileName ns2 searchKey
-        result2 `shouldBe` Nothing
+        queryEntry @TestEntry fileName ns2 searchKey `shouldReturn` Nothing
 
     it "handles many namespaces efficiently" $
       withSystemTempDirectory "query-test-XXXXXX" $ \tmpDir -> do
@@ -222,5 +211,4 @@ multiNamespaceTests =
 
         serializeTestData fileName namespaces
 
-        result <- queryEntry fileName searchNs searchKey
-        result `shouldBe` Just (mkTestEntry 5 1 (BS.pack [5]))
+        queryEntry fileName searchNs searchKey `shouldReturn` Just (mkTestEntry 5 1 (BS.pack [5]))

scls-util/src/Cardano/SCLS/Util/Check.hs

@@ -23,7 +23,7 @@ import Cardano.SCLS.CDDL (NamespaceInfo (..), namespaces)
 import Cardano.SCLS.Internal.Entry.CBOREntry (GenericCBOREntry (..))
 import Cardano.SCLS.Internal.Entry.ChunkEntry (ChunkEntry (..))
 import Cardano.SCLS.Internal.Hash (Digest (..), digest)
-import Cardano.SCLS.Internal.Reader (decodeChunkEntries, extractNamespaceList, withRecordData)
+import Cardano.SCLS.Internal.Reader (extractNamespaceList, streamChunkEntries, withRecordData)
 import Cardano.SCLS.Internal.Record.Chunk (Chunk (..))
 import Cardano.SCLS.Util.Result
 import Cardano.Types.Namespace (Namespace)
@@ -201,7 +201,7 @@ validateChunk cddlTrees Chunk{..} = do
       -- We do not known how to decode values inside, so we just read the data
       -- this way we can calculate count and check digest
       actualCount <-
-        decodeChunkEntries @(Entry RawBytes) chunkData & S.length_
+        streamChunkEntries @(Entry RawBytes) chunkData & S.length_
       pure $
         if actualCount /= fromIntegral chunkEntriesCount
           then [EntryCountMismatch chunkEntriesCount actualCount]
@@ -210,7 +210,7 @@ validateChunk cddlTrees Chunk{..} = do
       withSomeSNat keySize \(snat :: SNat n) -> do
         withKnownNat snat do
           (actualCount S.:> formatErrors S.:> ()) <-
-            decodeChunkEntries @(GenericCBOREntry n) chunkData
+            streamChunkEntries @(GenericCBOREntry n) chunkData
               & S.copy
               & S.zip (S.enumFrom 1)
               & S.mapMaybe (validateAgainst spec)