Fix MultiScan Prepare() to support dictionary compression (facebook#13896)

Summary:
I saw failure when added some asserts near https://github.com/facebook/rocksdb/blob/b9957c991cae44959f96888369caf1b145398132/table/block_based/block_based_table_iterator.cc#L1201-L1205 in stress test. The decompression failed with error message like "Corruption: Failed zlib inflate: -3". This PR fixes the issue to use the right decompressor for dictionary compression.

Pull Request resolved: facebook#13896

Test Plan: updated unit test that checks no I/O is done after Prepare(), this would fail before this change.

Reviewed By: anand1976

Differential Revision: D80821500

Pulled By: cbi42

fbshipit-source-id: a4322c0da99a2d10e9787d0ec168668567c0c19a

Author: cbi42

table/block_based/block_based_table_iterator.cc

@@ -1175,6 +1175,34 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
       }
     }
 
+    // Get compression dictionary if available - needed for dictionary-aware
+    // decompression
+    UnownedPtr<Decompressor> decompressor =
+        table_->get_rep()->decompressor.get();
+    CachableEntry<DecompressorDict> cached_dict;
+    if (table_->get_rep()->uncompression_dict_reader) {
+      s = table_->get_rep()
+              ->uncompression_dict_reader->GetOrReadUncompressionDictionary(
+                  /* prefetch_buffer= */ nullptr, read_options_,
+                  /* get_context= */ nullptr, /* lookup_context= */ nullptr,
+                  &cached_dict);
+      if (!s.ok()) {
+#ifndef NDEBUG
+        fprintf(stdout, "Prepare dictionary loading failed with %s\n",
+                s.ToString().c_str());
+#endif
+        // Abort: dictionary lookup failed.
+        return;
+      }
+      if (!cached_dict.GetValue()) {
+#ifndef NDEBUG
+        fprintf(stdout, "Success but no dictionary read\n");
+#endif
+        return;
+      }
+      decompressor = cached_dict.GetValue()->decompressor_.get();
+    }
+
     // Init blocks and pin them in block cache.
     MemoryAllocator* memory_allocator =
         table_->get_rep()->table_options.block_cache->memory_allocator();
@@ -1199,10 +1227,12 @@ void BlockBasedTableIterator::Prepare(const MultiScanArgs* multiscan_opts) {
 #endif
         assert(pinned_data_blocks_guard[block_idx].IsEmpty());
         s = table_->CreateAndPinBlockInCache<Block_kData>(
-            read_options_, block, table_->get_rep()->decompressor.get(),
-            &tmp_contents,
+            read_options_, block, decompressor, &tmp_contents,
             &(pinned_data_blocks_guard[block_idx].As<Block_kData>()));
         if (!s.ok()) {
+#ifndef NDEBUG
+          fprintf(stdout, "Prepare failed with %s\n", s.ToString().c_str());
+#endif
           // Abort: failed to create and pin block in cache
           return;
         }

table/block_based/block_based_table_reader_test.cc

@@ -994,6 +994,18 @@ TEST_P(BlockBasedTableReaderTestVerifyChecksum, ChecksumMismatch) {
 }
 
 TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
+  std::ostringstream param_trace;
+  param_trace << "[MultiScanPrepare] Test params: " << "CompressionType="
+              << CompressionTypeToString(compression_type_)
+              << ", UseDirectReads=" << (use_direct_reads_ ? "true" : "false")
+              << ", UDTEnabled=" << (udt_enabled_ ? "true" : "false")
+              << ", PersistUDT=" << (persist_udt_ ? "true" : "false")
+              << ", CompressionParallelThreads="
+              << compression_parallel_threads_
+              << ", CompressionDictBytes=" << compression_dict_bytes_
+              << ", SameKeyDiffTs=" << (same_key_diff_ts_ ? "true" : "false");
+  std::cout << param_trace.str() << std::endl;
+
   Options options;
   options.statistics = CreateDBStatistics();
   ReadOptions read_opts;
@@ -1052,6 +1064,9 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
     iter->Next();
   }
   ASSERT_OK(iter->status());
+  // No I/O expected during scanning since all blocks were loaded and pinned.
+  ASSERT_EQ(read_count_after,
+            options.statistics->getTickerCount(NON_LAST_LEVEL_READ_COUNT));
 
   iter.reset(table->NewIterator(
       read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
@@ -1089,7 +1104,7 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
       read_opts, options_.prefix_extractor.get(), /*arena=*/nullptr,
       /*skip_filters=*/false, TableReaderCaller::kUncategorized));
   // Should do two I/Os since blocks 80-81 and 90-95 are already in block cache,
-  // reads from blocks 50-79 and 82-.. are co
+  // reads from blocks 50-79 and 82-.. are coalesced.
   scan_options = MultiScanArgs(BytewiseComparator());
   scan_options.insert(ExtractUserKey(kv[50 * kEntriesPerBlock].first));
   read_count_before =
@@ -1106,6 +1121,8 @@ TEST_P(BlockBasedTableReaderTest, MultiScanPrepare) {
   }
   ASSERT_FALSE(iter->Valid());
   ASSERT_OK(iter->status());
+  ASSERT_EQ(read_count_after,
+            options.statistics->getTickerCount(NON_LAST_LEVEL_READ_COUNT));
 
   // Check cases when Seek key does not match start key in ScanOptions
   iter.reset(table->NewIterator(

unreleased_history/bug_fixes/multi-scan-dict-compression.md

@@ -0,0 +1 @@
+* Fix a bug in MultiScan that causes it to fall back to a normal scan when dictionary compression is enabled.