feat: Add cluster index support in velox selective reader (facebookincubator#406)

Summary:
Pull Request resolved: facebookincubator#406

This diff integrates the cluster index into the Nimble selective reader to enable filter-based row pruning at read time:

**NimbleRowReaderOptions**: Added format-specific options to control index-based filtering via `setIndexEnabled()`.

**ReaderBase refactoring**:
   - Changed constructor to static `create()` factory method for proper initialization ordering
   - Renamed `memoryPool_` to `pool_` for consistency
   - Added `StripeStreams::streamIndex()` to get stream index for decoders
   - Added `StripeStreams::enqueueKeyStream()` to load key stream data for IO coalescing with data streams

**SelectiveNimbleRowReader index integration**:
   - Added `convertIndexColumnsToFileSchema()` to convert index column names from nimble schema to file schema (user-facing names)
   - `initIndexBounds()`: Converts index columns to file schema names, then converts ScanSpec filters to encoded key bounds
   - `updateStartStripeFromLowerIndexBound()` / `updateEndStripeFromUpperIndexBound()`: Uses TabletIndex lookup to skip entire stripes outside the filter range
   - `buildIndexReader()`: Creates IndexReader for first/last stripes only (they can be the same) to find exact row positions
   - `setStripeRowRange()`: Adjusts row range within stripes based on key bounds, then calls `columnReader_->seekTo()` to skip to the starting position

**NimbleData updates**:
   - Changed `streams_` and `memoryPool_` from references to pointers for consistency
   - Updated decoder creation to pass stream index for index-aware decoding
   - Changed `VELOX_UNREACHABLE()` to `NIMBLE_UNSUPPORTED()` for unsupported operations

**Random Skip Handling with Cluster Index: When random sampling (random skip) is enabled along with cluster index filtering, the random skip tracker must be updated to account for rows that are skipped due to index bounds (not just rows that are actually read). This ensures that random sampling produces consistent results whether or not index filtering is enabled. The implementation handles three scenarios:
 - Leading stripes filtered out: When updateStartStripeFromLowerIndexBound() skips entire stripes at the beginning due to the lower bound, maybeUpdateRandomSkip() is called immediately with the total row count of those filtered stripes.
 - Trailing stripes filtered out: When updateEndStripeFromUpperIndexBound() skips entire stripes at the end due to the upper bound, the row counts are accumulated into trailingSkippedRows_.
Partial stripe rows filtered: When setStripeRowRange() adjusts the row range within stripes (rows skipped at the start of the first stripe due to lower bound, or rows skipped at the end of the last stripe due to upper bound), the skipped row counts are handled: leading rows trigger immediate maybeUpdateRandomSkip() call, while trailing rows are accumulated into trailingSkippedRows_.
- Deferred update for trailing rows: All trailing skipped rows (from both entire stripes and partial last stripe) are updated together via maybeUpdateRandomSkip(trailingSkippedRows_) in nextRowNumber() when reaching kAtEnd.
Added randomSkipWithIndex test case and integrated random skip testing into the randomSchemaAndFilters fuzzer test.

**E2EIndexTest**: Comprehensive test suite covering:
   - Single-column keys: bigint, double, float, timestamp, boolean, varchar
   - Multi-column composite keys with various filter combinations (point + point, point + range, range + point, etc.)
   - Fuzzer testing with randomly generated table and index schema and randomly generated index column filter for conversion and non-index column filter combinations.
   - Verification that index-enabled and index-disabled reads produce identical results

Reviewed By: Yuhta

Differential Revision: D90074133

fbshipit-source-id: b64323dc51303e3c5ab968e11d73f166e0a2f592

Author: xiaoxmeng

dwio/nimble/velox/selective/NimbleData.cpp

@@ -32,8 +32,8 @@ NimbleData::NimbleData(
         std::unique_ptr<Encoding>(velox::memory::MemoryPool&, std::string_view)>
         encodingFactory)
     : nimbleType_(nimbleType),
-      streams_(streams),
-      memoryPool_(memoryPool),
+      streams_(&streams),
+      pool_(&memoryPool),
       inMapDecoder_(inMapDecoder),
       encodingFactory_{encodingFactory} {
   switch (nimbleType->kind()) {
@@ -97,13 +97,13 @@ void NimbleData::readNulls(
   }
   auto numBytes = velox::bits::nbytes(numValues);
   if (!nulls || nulls->capacity() < numBytes) {
-    nulls = AlignedBuffer::allocate<char>(numBytes, &memoryPool_);
+    nulls = AlignedBuffer::allocate<char>(numBytes, pool_);
   }
   nulls->setSize(numBytes);
   auto* nullsPtr = nulls->asMutable<uint64_t>();
   if (inMapDecoder_) {
     if (nullsDecoder_) {
-      dwio::common::ensureCapacity<char>(inMap_, numBytes, &memoryPool_);
+      dwio::common::ensureCapacity<char>(inMap_, numBytes, pool_);
       inMapDecoder_->nextBools(
           inMap_->asMutable<uint64_t>(), numValues, incomingNulls);
       nullsDecoder_->nextBools(nullsPtr, numValues, inMap_->as<uint64_t>());
@@ -155,31 +155,34 @@ uint64_t NimbleData::skipNulls(uint64_t numValues, bool /*nullsOnly*/) {
 }
 
 ChunkedDecoder NimbleData::makeScalarDecoder() {
+  const auto streamId = nimbleType_->asScalar().scalarDescriptor().offset();
   return ChunkedDecoder(
-      streams_.enqueue(nimbleType_->asScalar().scalarDescriptor().offset()),
+      streams_->enqueue(streamId),
       /*decodeValuesWithNulls=*/false,
-      /*streamIndex=*/nullptr,
-      &memoryPool_);
+      streams_->streamIndex(streamId),
+      pool_);
 }
 
 ChunkedDecoder NimbleData::makeMicrosDecoder() {
   VELOX_CHECK(nimbleType_->isTimestampMicroNano());
+  const auto streamId =
+      nimbleType_->asTimestampMicroNano().microsDescriptor().offset();
   return ChunkedDecoder(
-      streams_.enqueue(
-          nimbleType_->asTimestampMicroNano().microsDescriptor().offset()),
+      streams_->enqueue(streamId),
       /*decodeValuesWithNulls=*/false,
-      /*streamIndex=*/nullptr,
-      &memoryPool_);
+      streams_->streamIndex(streamId),
+      pool_);
 }
 
 ChunkedDecoder NimbleData::makeNanosDecoder() {
   VELOX_CHECK(nimbleType_->isTimestampMicroNano());
+  const auto streamId =
+      nimbleType_->asTimestampMicroNano().nanosDescriptor().offset();
   return ChunkedDecoder(
-      streams_.enqueue(
-          nimbleType_->asTimestampMicroNano().nanosDescriptor().offset()),
+      streams_->enqueue(streamId),
       /*decodeValuesWithNulls=*/false,
-      /*streamIndex=*/nullptr,
-      &memoryPool_);
+      streams_->streamIndex(streamId),
+      pool_);
 }
 
 std::unique_ptr<ChunkedDecoder> NimbleData::makeLengthDecoder() {
@@ -199,16 +202,15 @@ std::unique_ptr<ChunkedDecoder> NimbleData::makeLengthDecoder() {
 std::unique_ptr<ChunkedDecoder> NimbleData::makeDecoder(
     const StreamDescriptor& descriptor,
     bool decodeValuesWithNulls) {
-  auto input = streams_.enqueue(descriptor.offset());
+  auto input = streams_->enqueue(descriptor.offset());
   if (!input) {
     return nullptr;
   }
   return std::make_unique<ChunkedDecoder>(
       std::move(input),
       decodeValuesWithNulls,
-      /*streamIndex=*/nullptr,
-      &memoryPool_,
-      encodingFactory_);
+      streams_->streamIndex(descriptor.offset()),
+      pool_);
 }
 
 std::unique_ptr<velox::dwio::common::FormatData> NimbleParams::toFormatData(

dwio/nimble/velox/selective/NimbleData.h

@@ -56,15 +56,15 @@ class NimbleData : public velox::dwio::common::FormatData {
 
   velox::dwio::common::PositionProvider seekToRowGroup(
       int64_t /*index*/) final {
-    VELOX_UNREACHABLE();
+    NIMBLE_UNSUPPORTED();
   }
 
   void filterRowGroups(
       const velox::common::ScanSpec& /*scanSpec*/,
       uint64_t /*rowsPerRowGroup*/,
       const velox::dwio::common::StatsContext& /*writerContext*/,
       FilterRowGroupsResult& /*result*/) final {
-    VELOX_UNREACHABLE();
+    NIMBLE_UNSUPPORTED();
   }
 
   const Type& nimbleType() const {
@@ -103,8 +103,8 @@ class NimbleData : public velox::dwio::common::FormatData {
       bool decodeValuesWithNulls);
 
   const std::shared_ptr<const Type> nimbleType_;
-  StripeStreams& streams_;
-  velox::memory::MemoryPool& memoryPool_;
+  StripeStreams* const streams_;
+  velox::memory::MemoryPool* const pool_;
   ChunkedDecoder* const inMapDecoder_;
   std::unique_ptr<ChunkedDecoder> nullsDecoder_;
   velox::BufferPtr inMap_;
@@ -169,7 +169,7 @@ class NimbleParams : public velox::dwio::common::FormatParams {
 
  private:
   const std::shared_ptr<const Type> nimbleType_;
-  StripeStreams* const streams_;
+  StripeStreams* const streams_{nullptr};
   RowSizeTracker* const rowSizeTracker_{nullptr};
   const bool preserveFlatMapsInMemory_{false};
   ChunkedDecoder* inMapDecoder_{nullptr};

dwio/nimble/velox/selective/NimbleRowReaderOptions.h

@@ -0,0 +1,46 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "velox/dwio/common/Options.h"
+
+namespace facebook::nimble {
+
+/// Format-specific options for Nimble row readers.
+class NimbleRowReaderOptions
+    : public velox::dwio::common::FormatSpecificOptions {
+ public:
+  NimbleRowReaderOptions() = default;
+
+  /// Sets whether to use the cluster index for filter-based row pruning.
+  /// When enabled (default), filters from ScanSpec are converted to index
+  /// bounds for efficient row skipping based on the file's cluster index.
+  /// When disabled, all rows are scanned without index-based pruning.
+  void setIndexEnabled(bool enabled) {
+    indexEnabled_ = enabled;
+  }
+
+  /// Returns true if cluster index-based filtering is enabled.
+  bool indexEnabled() const {
+    return indexEnabled_;
+  }
+
+ private:
+  bool indexEnabled_{true};
+};
+
+} // namespace facebook::nimble

dwio/nimble/velox/selective/ReaderBase.cpp

@@ -16,6 +16,7 @@
 
 #include "dwio/nimble/velox/selective/ReaderBase.h"
 
+#include "dwio/nimble/index/IndexConstants.h"
 #include "dwio/nimble/tablet/Constants.h"
 #include "dwio/nimble/velox/SchemaSerialization.h"
 #include "dwio/nimble/velox/SchemaUtils.h"
@@ -26,7 +27,6 @@ namespace facebook::nimble {
 using namespace facebook::velox;
 
 namespace {
-
 const std::string kSchemaSectionString(kSchemaSection);
 const std::vector<std::string> kPreloadOptionalSections = {
     kSchemaSectionString};
@@ -47,25 +47,50 @@ TypePtr getFileSchema(
   return dwio::common::Reader::updateColumnNames(
       fileSchema, options.fileSchema());
 }
-
 } // namespace
 
+std::shared_ptr<ReaderBase> ReaderBase::create(
+    std::unique_ptr<velox::dwio::common::BufferedInput> input,
+    const velox::dwio::common::ReaderOptions& options) {
+  // Initialize all members
+  auto tablet = TabletReader::create(
+      // TODO: Make TabletReader taking BufferedInput.
+      input->getReadFile().get(),
+      options.memoryPool(),
+      kPreloadOptionalSections);
+
+  auto* pool = &options.memoryPool();
+  const auto& randomSkip = options.randomSkip();
+  const auto& scanSpec = options.scanSpec();
+  const auto nimbleSchema = loadSchema(*tablet);
+  auto fileSchema =
+      asRowType(getFileSchema(options, convertToVeloxType(*nimbleSchema)));
+
+  return std::shared_ptr<ReaderBase>(new ReaderBase(
+      std::move(input),
+      std::move(tablet),
+      pool,
+      randomSkip,
+      scanSpec,
+      nimbleSchema,
+      std::move(fileSchema)));
+}
+
 ReaderBase::ReaderBase(
-    std::unique_ptr<dwio::common::BufferedInput> input,
-    const dwio::common::ReaderOptions& options)
-    : input_(std::move(input)),
-      tablet_(
-          TabletReader::create(
-              // TODO: Make TabletReader taking BufferedInput.
-              input_->getReadFile().get(),
-              options.memoryPool(),
-              kPreloadOptionalSections)),
-      memoryPool_(&options.memoryPool()),
-      randomSkip_(options.randomSkip()),
-      scanSpec_(options.scanSpec()),
-      nimbleSchema_(loadSchema(*tablet_)),
-      fileSchema_(asRowType(
-          getFileSchema(options, convertToVeloxType(*nimbleSchema_)))) {}
+    std::unique_ptr<velox::dwio::common::BufferedInput> input,
+    std::shared_ptr<TabletReader> tablet,
+    velox::memory::MemoryPool* pool,
+    const std::shared_ptr<velox::random::RandomSkipTracker>& randomSkip,
+    const std::shared_ptr<velox::common::ScanSpec>& scanSpec,
+    std::shared_ptr<const Type> nimbleSchema,
+    velox::RowTypePtr fileSchema)
+    : input_{std::move(input)},
+      tablet_{std::move(tablet)},
+      pool_{pool},
+      randomSkip_{randomSkip},
+      scanSpec_{scanSpec},
+      nimbleSchema_{std::move(nimbleSchema)},
+      fileSchema_{std::move(fileSchema)} {}
 
 std::optional<common::Region> StripeStreams::streamRegion(int streamId) const {
   NIMBLE_CHECK(stripeIdentifier_.has_value());
@@ -94,4 +119,27 @@ std::unique_ptr<dwio::common::SeekableInputStream> StripeStreams::enqueue(
   return readerBase_->input().enqueue(*region, &sid);
 }
 
+std::shared_ptr<index::StreamIndex> StripeStreams::streamIndex(
+    int streamId) const {
+  NIMBLE_CHECK(stripeIdentifier_.has_value());
+
+  const auto& indexGroup = stripeIdentifier_->indexGroup();
+  if (indexGroup == nullptr) {
+    return nullptr;
+  }
+  return indexGroup->createStreamIndex(stripe_, streamId);
+}
+
+std::unique_ptr<velox::dwio::common::SeekableInputStream>
+StripeStreams::enqueueKeyStream() {
+  NIMBLE_CHECK(stripeIdentifier_.has_value());
+
+  const auto& indexGroup = stripeIdentifier_->indexGroup();
+  NIMBLE_CHECK_NOT_NULL(indexGroup);
+
+  const auto region = indexGroup->keyStreamRegion(stripe_);
+  const dwio::common::StreamIdentifier sid(kKeyStreamId);
+  return readerBase_->input().enqueue(region, &sid);
+}
+
 } // namespace facebook::nimble

dwio/nimble/velox/selective/ReaderBase.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include "dwio/nimble/index/StripeIndexGroup.h"
 #include "dwio/nimble/tablet/TabletReader.h"
 #include "dwio/nimble/velox/SchemaReader.h"
 #include "velox/dwio/common/BufferedInput.h"
@@ -25,7 +26,7 @@ namespace facebook::nimble {
 
 class ReaderBase {
  public:
-  ReaderBase(
+  static std::shared_ptr<ReaderBase> create(
       std::unique_ptr<velox::dwio::common::BufferedInput> input,
       const velox::dwio::common::ReaderOptions& options);
 
@@ -37,8 +38,8 @@ class ReaderBase {
     return *tablet_;
   }
 
-  velox::memory::MemoryPool* memoryPool() const {
-    return memoryPool_;
+  velox::memory::MemoryPool* pool() const {
+    return pool_;
   }
 
   const std::shared_ptr<velox::random::RandomSkipTracker>& randomSkip() const {
@@ -64,9 +65,18 @@ class ReaderBase {
   }
 
  private:
+  ReaderBase(
+      std::unique_ptr<velox::dwio::common::BufferedInput> input,
+      std::shared_ptr<TabletReader> tablet,
+      velox::memory::MemoryPool* pool,
+      const std::shared_ptr<velox::random::RandomSkipTracker>& randomSkip,
+      const std::shared_ptr<velox::common::ScanSpec>& scanSpec,
+      std::shared_ptr<const Type> nimbleSchema,
+      velox::RowTypePtr fileSchema);
+
   const std::unique_ptr<velox::dwio::common::BufferedInput> input_;
   const std::shared_ptr<TabletReader> tablet_;
-  velox::memory::MemoryPool* const memoryPool_;
+  velox::memory::MemoryPool* const pool_;
   const std::shared_ptr<velox::random::RandomSkipTracker> randomSkip_;
   const std::shared_ptr<velox::common::ScanSpec> scanSpec_;
   const std::shared_ptr<const Type> nimbleSchema_;
@@ -80,9 +90,15 @@ class StripeStreams {
   explicit StripeStreams(const std::shared_ptr<ReaderBase>& readerBase)
       : readerBase_(readerBase) {}
 
-  void setStripe(int stripe) {
+  void setStripe(int stripe, bool loadIndex = false) {
     stripe_ = stripe;
-    stripeIdentifier_ = readerBase_->tablet().stripeIdentifier(stripe_);
+    stripeIdentifier_ =
+        readerBase_->tablet().stripeIdentifier(stripe_, loadIndex);
+    if (loadIndex) {
+      NIMBLE_CHECK_NOT_NULL(
+          stripeIdentifier_->indexGroup(),
+          "Index group should be set when loadIndex is true");
+    }
   }
 
   bool hasStream(int streamId) const {
@@ -92,10 +108,22 @@ class StripeStreams {
   std::unique_ptr<velox::dwio::common::SeekableInputStream> enqueue(
       int streamId);
 
+  std::unique_ptr<velox::dwio::common::SeekableInputStream> enqueueKeyStream();
+
   void load() {
     readerBase_->input().load(velox::dwio::common::LogType::STREAM_BUNDLE);
   }
 
+  int32_t stripeIndex() const {
+    return stripe_;
+  }
+
+  const std::shared_ptr<StripeIndexGroup>& indexGroup() const {
+    return stripeIdentifier_->indexGroup();
+  }
+
+  std::shared_ptr<index::StreamIndex> streamIndex(int streamId) const;
+
  private:
   std::optional<velox::common::Region> streamRegion(int streamId) const;