feat: Improve index processing and code refactor (#385)

Summary:
Pull Request resolved: #385

Move StripeIndexGroup and TabletIndex to facebook::nimble::index namespace for better organization
Add NIMBLE_USER_CHECK macro for user-facing validation errors
Improve error messages in EncodingFactory to include encoding type
Optimize StripePositionIndex by using accumulated chunk counts (prefix sums) instead of separate count and index arrays
Add IndexConfig.cpp implementation file and move defaultEncodingLayout out of header
Add test utilities (TestChunkDecoder, SingleChunkDecoder, StripeIndexGroupTestHelper) for chunk decoding and index testing
Add TabletReaderTestHelper for testing stripe/index group caching behavior

Reviewed By: HuamengJiang

Differential Revision: D89835718

fbshipit-source-id: 034d70dc66b3ac05629dadea4b1778656a0fa884

Author: xiaoxmeng

dwio/nimble/common/Exceptions.h

@@ -157,7 +157,10 @@ NIMBLE_DECLARE_CHECK_FAIL_TEMPLATES(::facebook::nimble::NimbleInternalError);
 
 #define NIMBLE_CHECK(expr, ...) _NIMBLE_CHECK_IMPL(expr, #expr, ##__VA_ARGS__)
 
-// Verify an expected file format conditions. Failure of this condition means
+#define NIMBLE_USER_CHECK(expr, ...) \
+  _NIMBLE_USER_CHECK_IMPL(expr, #expr, ##__VA_ARGS__)
+
+// Verify an expected file format conditions.
 // the file is corrupted (e.g. passed magic number and version verification, but
 // got unexpected format). This will trigger a user error.
 #define NIMBLE_CHECK_FILE(condition, ...)              \

dwio/nimble/encodings/EncodingFactory.cpp

@@ -190,7 +190,8 @@ std::unique_ptr<Encoding> EncodingFactory::decode(
     }
     default: {
       NIMBLE_UNREACHABLE(
-          "Trying to deserialize invalid EncodingType -- garbage input?");
+          "Trying to deserialize invalid EncodingType:{} -- garbage input?",
+          static_cast<int>(encodingType));
     }
   }
 }

dwio/nimble/index/CMakeLists.txt

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-add_library(nimble_index StripeIndexGroup.cpp TabletIndex.cpp)
+add_library(nimble_index IndexConfig.cpp StripeIndexGroup.cpp TabletIndex.cpp)
 target_link_libraries(
   nimble_index
   nimble_common

dwio/nimble/index/IndexConfig.cpp

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "dwio/nimble/index/IndexConfig.h"
+
+namespace facebook::nimble {
+
+EncodingLayout IndexConfig::defaultEncodingLayout(
+    CompressionType compressionType) {
+  return EncodingLayout{
+      EncodingType::Trivial,
+      compressionType,
+      {EncodingLayout{EncodingType::Trivial, CompressionType::Uncompressed}}};
+}
+
+} // namespace facebook::nimble

dwio/nimble/index/IndexConfig.h

@@ -50,12 +50,7 @@ struct IndexConfig {
   /// TrivialEncoding for string_view needs a nested encoding for string
   /// lengths.
   static EncodingLayout defaultEncodingLayout(
-      CompressionType compressionType = CompressionType::Zstd) {
-    return EncodingLayout{
-        EncodingType::Trivial,
-        compressionType,
-        {EncodingLayout{EncodingType::Trivial, CompressionType::Uncompressed}}};
-  }
+      CompressionType compressionType = CompressionType::Zstd);
 };
 
 } // namespace facebook::nimble

dwio/nimble/index/StripeIndexGroup.cpp

@@ -20,7 +20,7 @@
 
 #include <algorithm>
 
-namespace facebook::nimble {
+namespace facebook::nimble::index {
 
 namespace {
 template <typename T>
@@ -162,22 +162,23 @@ std::optional<ChunkLocation> StripeIndexGroup::lookupChunk(
   const auto* positionIndex = root->position_index();
   NIMBLE_CHECK_NOT_NULL(positionIndex);
 
-  // Use stream_chunk_counts and stream_chunk_indexes to find the range
+  // Use accumulated stream_chunk_counts to find the range
   // for this specific stream in this stripe
   const auto* streamChunkCounts = positionIndex->stream_chunk_counts();
   NIMBLE_CHECK_NOT_NULL(streamChunkCounts);
-  const auto* streamChunkIndexes = positionIndex->stream_chunk_indexes();
-  NIMBLE_CHECK_NOT_NULL(streamChunkIndexes);
-  NIMBLE_CHECK_EQ(streamChunkCounts->size(), streamChunkIndexes->size());
 
   // Calculate the index in the flattened stream_chunk_counts array
   const uint32_t streamChunkCountIndex = stripeOffset * streamCount_ + streamId;
   NIMBLE_CHECK_LT(streamChunkCountIndex, streamChunkCounts->size());
 
-  const uint32_t chunkCount = streamChunkCounts->Get(streamChunkCountIndex);
-  const uint32_t startChunkIndex =
-      streamChunkIndexes->Get(streamChunkCountIndex);
-  const uint32_t endChunkIndex = startChunkIndex + chunkCount;
+  // stream_chunk_counts is accumulated, so:
+  // - Start index: streamChunkCounts[streamChunkCountIndex - 1] (or 0 for
+  // first)
+  // - End index: streamChunkCounts[streamChunkCountIndex]
+  const uint32_t endChunkIndex = streamChunkCounts->Get(streamChunkCountIndex);
+  const uint32_t startChunkIndex = streamChunkCountIndex == 0
+      ? 0
+      : streamChunkCounts->Get(streamChunkCountIndex - 1);
 
   const auto* chunkRows = positionIndex->stream_chunk_rows();
   NIMBLE_CHECK_NOT_NULL(chunkRows);
@@ -252,4 +253,4 @@ std::optional<ChunkLocation> StreamIndex::lookupChunk(uint32_t rowId) const {
   return stripeIndexGroup_->lookupChunk(stripe_, streamId_, rowId);
 }
 
-} // namespace facebook::nimble
+} // namespace facebook::nimble::index

dwio/nimble/index/StripeIndexGroup.h

@@ -24,9 +24,14 @@
 #include "velox/common/file/Region.h"
 
 namespace facebook::nimble {
-
 class MetadataBuffer;
-class StripeGroup;
+} // namespace facebook::nimble
+
+namespace facebook::nimble::index {
+
+namespace test {
+class StripeIndexGroupTestHelper;
+} // namespace test
 
 /// Represents the location of a chunk within a stream.
 /// Both offsets are relative:
@@ -96,6 +101,8 @@ class StripeIndexGroup {
   const uint32_t firstStripe_;
   const uint32_t stripeCount_;
   const uint32_t streamCount_;
+
+  friend class test::StripeIndexGroupTestHelper;
 };
 
 /// StreamIndex is a convenience wrapper around StripeIndexGroup for rowId-based
@@ -122,4 +129,4 @@ class StreamIndex {
   const uint32_t streamId_;
 };
 
-} // namespace facebook::nimble
+} // namespace facebook::nimble::index

dwio/nimble/index/TabletIndex.cpp

@@ -21,7 +21,7 @@
 #include "dwio/nimble/tablet/IndexGenerated.h"
 #include "dwio/nimble/tablet/MetadataBuffer.h"
 
-namespace facebook::nimble {
+namespace facebook::nimble::index {
 
 namespace {
 const serialization::Index* getIndexRoot(const Section& indexSection) {
@@ -38,6 +38,13 @@ uint32_t getStripeCount(const Section& indexSection) {
   return stripeKeys->size() - 1;
 }
 
+uint32_t getIndexGroupCount(const Section& indexSection) {
+  const auto* indexRoot = getIndexRoot(indexSection);
+  const auto* stripeIndexGroups = indexRoot->stripe_index_groups();
+  NIMBLE_CHECK_NOT_NULL(stripeIndexGroups);
+  return stripeIndexGroups->size();
+}
+
 std::vector<std::string_view> getStripeKeys(const Section& indexSection) {
   const auto* indexRoot = getIndexRoot(indexSection);
   const auto* stripeKeys = indexRoot->stripe_keys();
@@ -71,18 +78,16 @@ std::unique_ptr<TabletIndex> TabletIndex::create(Section indexSection) {
 TabletIndex::TabletIndex(Section indexSection)
     : indexSection_{std::move(indexSection)},
       numStripes_{getStripeCount(indexSection_)},
+      numIndexGroups_{getIndexGroupCount(indexSection_)},
       stripeKeys_{getStripeKeys(indexSection_)},
       indexColumns_{getIndexColumns(indexSection_)} {
   NIMBLE_CHECK(!indexColumns_.empty());
   NIMBLE_CHECK_EQ(numStripes_ + 1, stripeKeys_.size());
 
   // Validate consistency between stripe group indices and stripe index groups
-  const auto* indexRoot = getIndexRoot(indexSection_);
-  const auto* stripeIndexGroups = indexRoot->stripe_index_groups();
-  NIMBLE_CHECK_NOT_NULL(stripeIndexGroups);
   NIMBLE_CHECK_GE(
       numStripes_,
-      stripeIndexGroups->size(),
+      numIndexGroups_,
       "Number of stripe index groups cannot exceed number of stripes");
 }
 
@@ -113,10 +118,10 @@ std::optional<StripeLocation> TabletIndex::lookup(
 }
 
 MetadataSection TabletIndex::groupIndexMetadata(uint32_t groupIndex) const {
+  NIMBLE_CHECK_LT(groupIndex, numIndexGroups_);
   const auto* indexRoot = getIndexRoot(indexSection_);
   const auto* indexGroupSections = indexRoot->stripe_index_groups();
   NIMBLE_CHECK_NOT_NULL(indexGroupSections);
-  NIMBLE_CHECK_LT(groupIndex, indexGroupSections->size());
   const auto* metadata = indexGroupSections->Get(groupIndex);
   return MetadataSection{
       metadata->offset(),
@@ -127,4 +132,4 @@ MetadataSection TabletIndex::groupIndexMetadata(uint32_t groupIndex) const {
 const MetadataBuffer& TabletIndex::rootIndex() const {
   return indexSection_.buffer();
 }
-} // namespace facebook::nimble
+} // namespace facebook::nimble::index

dwio/nimble/index/TabletIndex.h

@@ -22,7 +22,7 @@
 
 #include "dwio/nimble/tablet/MetadataBuffer.h"
 
-namespace facebook::nimble {
+namespace facebook::nimble::index {
 
 /// Represents the global stripe index within a tablet.
 struct StripeLocation {
@@ -126,13 +126,19 @@ class TabletIndex {
   /// This is needed by StripeIndexGroup to read stripe group information.
   const MetadataBuffer& rootIndex() const;
 
+  /// Returns the number of index groups for testing purposes.
+  size_t numIndexGroups() const {
+    return numIndexGroups_;
+  }
+
  private:
   explicit TabletIndex(Section indexSection);
 
   const Section indexSection_;
   const uint32_t numStripes_;
+  const uint32_t numIndexGroups_;
   const std::vector<std::string_view> stripeKeys_;
   const std::vector<std::string> indexColumns_;
 };
 
-} // namespace facebook::nimble
+} // namespace facebook::nimble::index

dwio/nimble/index/tests/StripeIndexGroupTest.cpp

@@ -22,7 +22,7 @@
 #include "dwio/nimble/index/tests/TabletIndexTestBase.h"
 #include "dwio/nimble/tablet/TabletReader.h"
 
-namespace facebook::nimble::test {
+namespace facebook::nimble::index::test {
 
 class StripeIndexGroupTest : public TabletIndexTestBase {};
 
@@ -664,4 +664,4 @@ TEST_F(StripeIndexGroupTest, stripeIndexAndStreamIdOutOfBound) {
   NIMBLE_ASSERT_THROW(
       stripeIndexGroup1->lookupChunk(3, 5, 0), "streamId out of range");
 }
-} // namespace facebook::nimble::test
+} // namespace facebook::nimble::index::test