feat: Add fast path to PrefixEncoding when no duplicates (facebookincubator#474)

Summary:
Pull Request resolved: facebookincubator#474

X-link: facebookincubator/velox#16321

The change allows PrefixEncoding to use noDuplicate parameter to go for a fast path searching (terminate early without needing to continuing binary search or scanning).

Reviewed By: xiaoxmeng

Differential Revision: D92676762

fbshipit-source-id: bda5412f54f13779acbfbfa8ec21f188a107d244

Author: tanjialiang

CMakeLists.txt

@@ -114,6 +114,9 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${SCRIPT_CXX_FLAGS}")
 
 message("FINAL CMAKE_CXX_FLAGS=${CMAKE_CXX_FLAGS}")
 
+add_compile_options(-ffunction-sections -fdata-sections)
+add_link_options(-Wl,--gc-sections)
+
 include(CTest) # include after project() but before add_subdirectory()
 
 find_package(flatbuffers)

dwio/nimble/encodings/Encoding.h

@@ -134,6 +134,11 @@ class Encoding {
   // iterator; if you need to move your row pointer back, reset() and skip().
   virtual void skip(uint32_t rowCount) = 0;
 
+  /// Seeks to the position at or after the given value.
+  ///
+  /// @param value Pointer to target value to seek.
+  /// @return Row index if found, std::nullopt if value is greater than all
+  /// entries.
   virtual std::optional<uint32_t> seekAtOrAfter(const void* value) {
     NIMBLE_UNSUPPORTED("seekAtOrAfter is not supported.");
   }

dwio/nimble/encodings/PrefixEncoding.cpp

@@ -163,8 +163,6 @@ std::optional<uint32_t> PrefixEncoding::seekAtOrAfter(const void* value) {
   const auto& targetValue = *static_cast<const std::string_view*>(value);
 
   // Binary search among restart points to find the block containing the target.
-  // Since keys are unique and sorted, we find the last restart point whose
-  // value is <= targetValue.
   uint32_t left = 0;
   uint32_t right = numRestarts_;
 
@@ -173,12 +171,8 @@ std::optional<uint32_t> PrefixEncoding::seekAtOrAfter(const void* value) {
 
     seekToRestartPoint(mid);
     const std::string_view restartValue = decodeEntry();
-    if (restartValue == targetValue) {
-      // Exact match found at restart point. Since keys are unique (no
-      // duplicates), we can return immediately.
-      return currentRow_ - 1;
-    }
-    if (restartValue < targetValue) {
+
+    if (restartValue.compare(targetValue) < 0) {
       left = mid + 1;
     } else {
       right = mid;
@@ -199,9 +193,6 @@ std::optional<uint32_t> PrefixEncoding::seekAtOrAfter(const void* value) {
   // Linear scan within the block
   while (currentRow_ < rowCount_) {
     const std::string_view currentValue = decodeEntry();
-    // Return when we find either:
-    // 1. Exact match (since keys are unique, no need to scan further), or
-    // 2. First value greater than target.
     if (currentValue >= targetValue) {
       return currentRow_ - 1;
     }

dwio/nimble/encodings/tests/PrefixEncodingTest.cpp

@@ -1157,4 +1157,175 @@ TEST_F(PrefixEncodingTest, fuzzerEncode) {
   }
 }
 
+// Test seekAtOrAfter with duplicate keys that span across restart intervals.
+// This specifically tests the fix where seeking for a key that matches
+// a restart point value should return the earliest occurrence, which may be
+// in the previous restart interval.
+TEST_F(PrefixEncodingTest, seekExactMatchWithDuplicatesAcrossRestarts) {
+  // Test case: duplicates at restart boundary
+  // Restart interval is 16, so restart points are at indices 0, 16, 32, etc.
+  // We create data where duplicate keys span the restart boundary.
+  struct TestCase {
+    std::string_view name;
+    std::vector<std::string> values;
+    std::string seekKey;
+    uint32_t expectedPosition;
+  };
+
+  const std::vector<TestCase> testCases = {
+      // Duplicate key "key_16" appears at indices 14, 15, 16, 17
+      // Restart point at index 16 has value "key_16"
+      // Seeking for "key_16" should return 14 (first occurrence in previous
+      // interval)
+      {"duplicates before restart point (index 16)",
+       // Indices 0-13: unique keys, 14-17: "key_16", 18-31: unique keys
+       []() {
+         std::vector<std::string> v;
+         v.reserve(14);
+         for (int i = 0; i < 14; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         // Duplicates spanning restart boundary (indices 14, 15, 16, 17)
+         for (int i = 0; i < 4; ++i) {
+           v.emplace_back("key_16");
+         }
+         for (int i = 18; i < 32; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         return v;
+       }(),
+       "key_16",
+       14},
+
+      // Duplicate key at restart point with more duplicates before
+      // Restart point at index 32, duplicates at indices 28-35
+      {"duplicates before restart point (index 32)",
+       []() {
+         std::vector<std::string> v;
+         v.reserve(28);
+         for (int i = 0; i < 28; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         // Duplicates spanning restart boundary at index 32 (indices 28-35)
+         for (int i = 0; i < 8; ++i) {
+           v.emplace_back("key_32");
+         }
+         for (int i = 36; i < 50; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         return v;
+       }(),
+       "key_32",
+       28},
+
+      // Single duplicate before restart point
+      {"single duplicate before restart point",
+       []() {
+         std::vector<std::string> v;
+         v.reserve(15);
+         for (int i = 0; i < 15; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         // Index 15 and 16 both have "key_16"
+         v.emplace_back("key_16");
+         v.emplace_back("key_16");
+         for (int i = 17; i < 32; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         return v;
+       }(),
+       "key_16",
+       15},
+
+      // Duplicates spanning multiple restart intervals
+      // Restart points at 16, 32; duplicates from index 14 to 34
+      {"duplicates spanning multiple restart intervals",
+       []() {
+         std::vector<std::string> v;
+         v.reserve(14);
+         for (int i = 0; i < 14; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         // Duplicates from index 14 to 34 (spanning restarts at 16 and 32)
+         for (int i = 0; i < 21; ++i) {
+           v.emplace_back("key_dup");
+         }
+         for (int i = 35; i < 50; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         return v;
+       }(),
+       "key_dup",
+       14},
+
+      // Duplicates at restart point but none before (should still work)
+      {"duplicates at restart point only",
+       []() {
+         std::vector<std::string> v;
+         v.reserve(16);
+         for (int i = 0; i < 16; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         // Duplicates starting exactly at restart point (indices 16-19)
+         for (int i = 0; i < 4; ++i) {
+           v.emplace_back("key_20");
+         }
+         for (int i = 20; i < 32; ++i) {
+           v.push_back(fmt::format("key_{:02d}", i));
+         }
+         return v;
+       }(),
+       "key_20",
+       16},
+
+      // All same values (extreme case)
+      {"all same values",
+       []() {
+         std::vector<std::string> v;
+         v.reserve(50);
+         for (int i = 0; i < 50; ++i) {
+           v.emplace_back("same_key");
+         }
+         return v;
+       }(),
+       "same_key",
+       0},
+  };
+
+  for (const auto& testCase : testCases) {
+    SCOPED_TRACE(testCase.name);
+
+    std::vector<std::string_view> values;
+    values.reserve(testCase.values.size());
+    for (const auto& s : testCase.values) {
+      values.push_back(s);
+    }
+
+    Buffer buffer{*pool_};
+    auto encoded = EncodingFactory::encode<std::string_view>(
+        createSelectionPolicy(), values, buffer);
+    stringBuffers_.clear();
+    auto encoding =
+        EncodingFactory::decode(*pool_, encoded, createStringBufferFactory());
+
+    std::string_view seekKey = testCase.seekKey;
+    auto result = encoding->seekAtOrAfter(&seekKey);
+    ASSERT_TRUE(result.has_value())
+        << "Expected to find key: " << testCase.seekKey;
+    EXPECT_EQ(result.value(), testCase.expectedPosition)
+        << "Expected position " << testCase.expectedPosition << " for key "
+        << testCase.seekKey << " but got " << result.value();
+
+    // Also verify that the value at the returned position matches
+    encoding->reset();
+    if (result.value() > 0) {
+      encoding->skip(result.value());
+    }
+    std::string_view decoded;
+    encoding->materialize(1, &decoded);
+    EXPECT_EQ(decoded, testCase.seekKey)
+        << "Value at position " << result.value() << " should be "
+        << testCase.seekKey;
+  }
+}
 } // namespace facebook::nimble::test

dwio/nimble/encodings/tests/TrivialEncodingTest.cpp

@@ -156,6 +156,91 @@ TEST_P(TrivialEncodingTestWithCompression, seekExactMatch) {
   }
 }
 
+TEST_P(TrivialEncodingTestWithCompression, seekExactMatchWithDuplicateValues) {
+  const auto values = toVector(
+      {"apple", "apple", "banana", "banana", "banana", "cherry", "cherry"});
+  const auto encoding = createEncoding(values, GetParam());
+
+  struct {
+    std::string_view seekValue;
+    std::optional<uint32_t> expected;
+
+    std::string debugString() const {
+      if (expected.has_value()) {
+        return fmt::format("seek '{}' -> row {}", seekValue, expected.value());
+      }
+      return fmt::format("seek '{}' -> not found", seekValue);
+    }
+  } testCases[] = {
+      // Should return first occurrence of duplicate values
+      {"apple", 0},
+      {"banana", 2},
+      {"cherry", 5},
+      // Between values (should return first of next value)
+      {"aaa", 0},
+      {"app", 0},
+      {"az", 2},
+      {"car", 5},
+      // After all values (not found)
+      {"date", std::nullopt},
+      {"zebra", std::nullopt}};
+
+  for (const auto& testCase : testCases) {
+    SCOPED_TRACE(testCase.debugString());
+    encoding->reset();
+    auto result = encoding->seekAtOrAfter(&testCase.seekValue);
+    ASSERT_EQ(result.has_value(), testCase.expected.has_value());
+    if (testCase.expected.has_value()) {
+      EXPECT_EQ(result.value(), testCase.expected.value());
+    }
+  }
+}
+
+TEST_P(TrivialEncodingTestWithCompression, seekAfterValueWithDuplicateValues) {
+  auto compressionType = GetParam();
+  // Duplicates at beginning (apple), middle (cherry), and end (kiwi)
+  auto values = toVector(
+      {"apple", "apple", "apple", "cherry", "cherry", "grape", "kiwi", "kiwi"});
+  auto encoding = createEncoding(values, compressionType);
+
+  struct {
+    std::string_view seekValue;
+    std::optional<uint32_t> expected;
+
+    std::string debugString() const {
+      if (expected.has_value()) {
+        return fmt::format("seek '{}' -> row {}", seekValue, expected.value());
+      }
+      return fmt::format("seek '{}' -> not found", seekValue);
+    }
+  } testCases[] = {
+      // Before first duplicate group
+      {"", 0},
+      {"aaa", 0},
+      // Between apple (end) and cherry (start)
+      {"banana", 3},
+      {"car", 3},
+      // Between cherry (end) and grape
+      {"date", 5},
+      {"fig", 5},
+      // Between grape and kiwi (start of end duplicates)
+      {"honey", 6},
+      // After last duplicate group (not found)
+      {"lemon", std::nullopt},
+      {"zebra", std::nullopt},
+  };
+
+  for (const auto& testCase : testCases) {
+    SCOPED_TRACE(testCase.debugString());
+    encoding->reset();
+    auto result = encoding->seekAtOrAfter(&testCase.seekValue);
+    ASSERT_EQ(result.has_value(), testCase.expected.has_value());
+    if (testCase.expected.has_value()) {
+      EXPECT_EQ(result.value(), testCase.expected.value());
+    }
+  }
+}
+
 TEST_P(TrivialEncodingTestWithCompression, seekExactMatchWithEmptyEncoding) {
   const auto values = toVector({});
   const auto encoding = createEncoding(values, GetParam());

dwio/nimble/index/IndexReader.cpp

@@ -35,11 +35,9 @@ IndexReader::IndexReader(
     std::unique_ptr<velox::dwio::common::SeekableInputStream> input,
     uint32_t stripeIndex,
     std::shared_ptr<StripeIndexGroup> indexGroup,
-    bool noDuplicateKey,
     velox::memory::MemoryPool* pool)
     : stripeIndex_{stripeIndex},
       indexGroup_{std::move(indexGroup)},
-      noDuplicateKey_{noDuplicateKey},
       input_{std::move(input)},
       pool_{pool} {
   NIMBLE_CHECK_NOT_NULL(input_);
@@ -50,14 +48,9 @@ std::unique_ptr<IndexReader> IndexReader::create(
     std::unique_ptr<velox::dwio::common::SeekableInputStream> input,
     uint32_t stripeIndex,
     std::shared_ptr<StripeIndexGroup> indexGroup,
-    bool noDuplicateKey,
     velox::memory::MemoryPool* pool) {
   return std::unique_ptr<IndexReader>(new IndexReader(
-      std::move(input),
-      stripeIndex,
-      std::move(indexGroup),
-      noDuplicateKey,
-      pool));
+      std::move(input), stripeIndex, std::move(indexGroup), pool));
 }
 
 std::optional<uint32_t> IndexReader::seekAtOrAfter(

dwio/nimble/index/IndexReader.h

@@ -44,7 +44,6 @@ class IndexReader {
       std::unique_ptr<velox::dwio::common::SeekableInputStream> input,
       uint32_t stripeIndex,
       std::shared_ptr<StripeIndexGroup> indexGroup,
-      bool noDuplicateKey,
       velox::memory::MemoryPool* pool);
 
   /// Seek to the position at or after the given encoded key.
@@ -57,7 +56,6 @@ class IndexReader {
       std::unique_ptr<velox::dwio::common::SeekableInputStream> input,
       uint32_t stripeIndex,
       std::shared_ptr<StripeIndexGroup> indexGroup,
-      bool noDuplicateKey,
       velox::memory::MemoryPool* pool);
 
   /// Seek to a key at or after the given encoded key within a chunk.
@@ -80,7 +78,6 @@ class IndexReader {
 
   const uint32_t stripeIndex_;
   const std::shared_ptr<StripeIndexGroup> indexGroup_;
-  const bool noDuplicateKey_;
 
   const std::unique_ptr<velox::dwio::common::SeekableInputStream> input_;
   velox::memory::MemoryPool* const pool_;

dwio/nimble/index/TabletIndex.cpp

@@ -80,10 +80,6 @@ std::vector<SortOrder> getSortOrders(const serialization::Index* indexRoot) {
   return result;
 }
 
-bool getNoDuplicateKey(const serialization::Index* indexRoot) {
-  return indexRoot->no_duplicate_key();
-}
-
 } // namespace
 
 std::unique_ptr<TabletIndex> TabletIndex::create(Section indexSection) {
@@ -97,8 +93,7 @@ TabletIndex::TabletIndex(Section indexSection)
       numIndexGroups_{getIndexGroupCount(indexRoot_)},
       stripeKeys_{getStripeKeys(indexRoot_)},
       indexColumns_{getIndexColumns(indexRoot_)},
-      sortOrders_{getSortOrders(indexRoot_)},
-      noDuplicateKey_{getNoDuplicateKey(indexRoot_)} {
+      sortOrders_{getSortOrders(indexRoot_)} {
   NIMBLE_CHECK(!indexColumns_.empty());
   NIMBLE_CHECK_EQ(numStripes_ + 1, stripeKeys_.size());
   NIMBLE_CHECK_EQ(indexColumns_.size(), sortOrders_.size());