feat: Add IndexReader interface for index-based lookups (#16330)

Summary:
Pull Request resolved: #16330

This diff introduces the `IndexReader` interface in `Reader.h` to provide a clean abstraction for index-based data lookups, separate from the general-purpose `RowReader` interface.

**New `IndexReader` class** - Abstract interface for index-based lookups with methods:
   - `encodeIndexBounds()` - Encodes index bounds into format-specific encoded key bounds
   - `lookupStripes()` - Looks up stripes that contain data matching the encoded key bounds
   - `setStripeRowRanges()` - Sets up row ranges for reading a specific stripe based on encoded bounds
   - `next()` - Pure virtual method to fetch the next portion of rows (without mutation support)

**New `createIndexReader()` method in `Reader`** - Factory method to create an `IndexReader` instance. Default implementation throws `VELOX_UNSUPPORTED`, allowing format-specific readers (e.g., Nimble) to override and provide index reading support.

**Supporting data structures** - `RowRange`, `StripeRowRanges`, and `StripeLookupResult` structs for representing row ranges and stripe lookup results.

This separation allows `HiveIndexReader` to use a dedicated `IndexReader` interface optimized for batched key-based lookups, while keeping the `RowReader` interface focused on sequential row-by-row reading with mutation support.

Reviewed By: Yuhta

Differential Revision: D92851481

fbshipit-source-id: 09838895db5a7dfe895a4722fa7d9223aeccc49d

Author: xiaoxmeng

velox/dwio/common/Reader.h

@@ -27,6 +27,7 @@
 #include "velox/dwio/common/SelectiveColumnReader.h"
 #include "velox/dwio/common/Statistics.h"
 #include "velox/dwio/common/TypeWithId.h"
+#include "velox/serializers/KeyEncoder.h"
 #include "velox/type/Type.h"
 #include "velox/vector/BaseVector.h"
 
@@ -186,6 +187,92 @@ class RowReader {
       VectorPtr& result);
 };
 
+/// Represents a row range within a stripe [startRow, endRow).
+struct RowRange {
+  vector_size_t startRow{0}; // Inclusive
+  vector_size_t endRow{0}; // Exclusive
+
+  RowRange() = default;
+  RowRange(vector_size_t _startRow, vector_size_t _endRow)
+      : startRow(_startRow), endRow(_endRow) {}
+
+  /// Returns true if this row range is empty (no rows to read).
+  bool empty() const {
+    return startRow >= endRow;
+  }
+};
+
+/**
+ * Abstract index reader interface for index-based lookups.
+ *
+ * IndexReader provides methods for encoding index bounds, looking up stripes,
+ * and reading data within specific row ranges. This interface is used by
+ * HiveIndexReader to perform efficient key-based lookups on indexed files.
+ */
+class IndexReader {
+ public:
+  virtual ~IndexReader() = default;
+
+  using KeyBoundsVector = std::vector<velox::serializer::EncodedKeyBounds>;
+
+  /// Encodes index bounds into format-specific encoded key bounds.
+  /// Different file formats may use different key encoding schemes, so this
+  /// allows the format-specific reader to handle the encoding.
+  ///
+  /// @param indexBounds The index bounds to encode, containing column names
+  ///        and lower/upper bound values.
+  /// @return A vector of encoded key bounds, one for each row in the input
+  ///         bounds.
+  /// @throws if encoding is not supported by the implementation or if any
+  ///         index bound fails to encode.
+  virtual KeyBoundsVector encodeIndexBounds(
+      const velox::serializer::IndexBounds& indexBounds) = 0;
+
+  /// Looks up stripes that contain data matching the encoded key bounds.
+  /// For each request row, returns the list of stripe indices that may contain
+  /// matching data based on the encoded lower and upper key bounds.
+  ///
+  /// @param keyBounds The encoded key bounds for each request row.
+  /// @return Stripe indices for each request row. Each inner vector contains
+  ///         the indices of stripes that may contain matching data for that
+  ///         request.
+  /// @throws if lookup is not supported by the implementation.
+  virtual std::vector<std::vector<uint32_t>> lookupStripes(
+      const KeyBoundsVector& keyBounds) = 0;
+
+  /// Looks up row ranges within a specific stripe based on encoded key bounds.
+  /// Computes row ranges per request without setting up state for iteration.
+  ///
+  /// @param stripeIndex The index of the stripe to compute row ranges for.
+  /// @param keyBounds The encoded key bounds for each request.
+  /// @return Row ranges for each request, one per input encoded key bounds.
+  ///         Empty ranges (startRow >= endRow) are included for requests with
+  ///         no matching data.
+  /// @throws if lookup is not supported by the implementation.
+  virtual std::vector<RowRange> lookupRowRanges(
+      uint32_t stripeIndex,
+      const KeyBoundsVector& keyBounds) = 0;
+
+  /// Sets row ranges for reading from a specific stripe. Must be called before
+  /// next() to set up the iteration state.
+  ///
+  /// @param stripeIndex The index of the stripe to read from.
+  /// @param rowRanges The row ranges to read within the stripe.
+  /// @throws if setting row ranges is not supported by the implementation.
+  virtual void setRowRanges(
+      uint32_t stripeIndex,
+      const std::vector<RowRange>& rowRanges) = 0;
+
+  /**
+   * Fetch the next portion of rows.
+   * @param size Max number of rows to read
+   * @param result output vector
+   * @return number of rows scanned in the file (including any rows filtered
+   * out), 0 if there are no more rows to read.
+   */
+  virtual uint64_t next(uint64_t size, velox::VectorPtr& result) = 0;
+};
+
 /**
  * Abstract reader class.
  *
@@ -236,6 +323,17 @@ class Reader {
   virtual std::unique_ptr<RowReader> createRowReader(
       const RowReaderOptions& options = {}) const = 0;
 
+  /**
+   * Create index reader object for index-based lookups.
+   * @param options Row reader options describing the data to fetch
+   * @return Index reader for efficient key-based lookups
+   * @throws if index reading is not supported by the implementation
+   */
+  virtual std::unique_ptr<IndexReader> createIndexReader(
+      const RowReaderOptions& options = {}) const {
+    VELOX_UNSUPPORTED("Reader::createIndexReader() is not supported");
+  }
+
   static TypePtr updateColumnNames(
       const TypePtr& fileType,
       const TypePtr& tableType);

velox/dwio/common/tests/ReaderTest.cpp

@@ -211,5 +211,17 @@ TEST_F(ReaderTest, projectColumnsMutation) {
   EXPECT_NE(0, numNonMax);
 }
 
+TEST_F(ReaderTest, rowRangeEmpty) {
+  // Empty when startRow >= endRow
+  EXPECT_TRUE((RowRange{0, 0}.empty()));
+  EXPECT_TRUE((RowRange{5, 5}.empty()));
+  EXPECT_TRUE((RowRange{10, 5}.empty()));
+
+  // Not empty when startRow < endRow
+  EXPECT_FALSE((RowRange{0, 1}.empty()));
+  EXPECT_FALSE((RowRange{0, 10}.empty()));
+  EXPECT_FALSE((RowRange{5, 10}.empty()));
+}
+
 } // namespace
 } // namespace facebook::velox::dwio::common

velox/serializers/KeyEncoder.cpp

@@ -20,34 +20,41 @@
 #include "velox/vector/FlatVector.h"
 
 namespace facebook::velox::serializer {
+namespace {
 
-bool IndexBounds::validate() const {
-  if (!lowerBound.has_value() && !upperBound.has_value()) {
+bool validateBound(
+    const IndexBound& indexBound,
+    const std::vector<std::string>& indexColumns) {
+  if (indexBound.bound == nullptr || indexBound.bound->size() == 0) {
     return false;
   }
 
-  const auto validateBound = [this](const IndexBound& bound) {
-    if (bound.bound == nullptr || bound.bound->size() == 0) {
+  const auto& rowType = asRowType(indexBound.bound->type());
+  if (rowType->size() != indexColumns.size()) {
+    return false;
+  }
+  for (const auto& columnName : indexColumns) {
+    if (!rowType->containsChild(columnName)) {
       return false;
     }
+  }
+  return true;
+}
 
-    const auto& rowType = asRowType(bound.bound->type());
-    if (rowType->size() != indexColumns.size()) {
-      return false;
-    }
-    for (const auto& columnName : indexColumns) {
-      if (!rowType->containsChild(columnName)) {
-        return false;
-      }
-    }
-    return true;
-  };
+} // namespace
+
+bool IndexBounds::validate() const {
+  if (!lowerBound.has_value() && !upperBound.has_value()) {
+    return false;
+  }
 
-  if (lowerBound.has_value() && !validateBound(lowerBound.value())) {
+  if (lowerBound.has_value() &&
+      !validateBound(lowerBound.value(), indexColumns)) {
     return false;
   }
 
-  if (upperBound.has_value() && !validateBound(upperBound.value())) {
+  if (upperBound.has_value() &&
+      !validateBound(upperBound.value(), indexColumns)) {
     return false;
   }
 

velox/serializers/KeyEncoder.h

@@ -174,14 +174,12 @@ class KeyEncoder {
   /// Increment fails when values are at their maximum (e.g., INT_MAX, strings
   /// with all \xFF characters, or nulls in NULLS_LAST ordering).
   ///
-  /// For multi-row bounds, returns a vector with one EncodedKeyBounds per row.
-  /// Each row is processed independently.
-  /// Encodes index bounds into byte-comparable key strings.
   /// Takes an IndexBounds containing lower and/or upper bounds and encodes them
-  /// into EncodedKeyBounds for efficient range comparison.
-  /// Throws if any lower bound fails to bump up (for exclusive bounds).
-  /// For upper bound bump up failures, the upperKey is set to std::nullopt
-  /// (unbounded).
+  /// into EncodedKeyBounds for efficient range comparison. Returns a vector
+  /// with one EncodedKeyBounds per row in 'indexBounds'. Each row is encoded
+  /// into a byte-comparable key string. Throws if any lower bound fails to bump
+  /// up (for exclusive bounds). For upper bound bump up failures, the upperKey
+  /// is set to std::nullopt (unbounded).
   std::vector<EncodedKeyBounds> encodeIndexBounds(
       const IndexBounds& indexBounds);
 
@@ -204,8 +202,10 @@ class KeyEncoder {
   std::vector<std::string> encode(const RowVectorPtr& input);
 
   // Creates a new row vector with the key columns incremented by 1 for multiple
-  // rows. Returns nullptr if any row fails to increment (all key columns
-  // overflow), otherwise returns RowVectorPtr with incremented values.
+  // rows. For each row, the increment takes place from the rightmost (least
+  // significant) column.
+  // Returns nullptr if any row fails to increment (all key columns overflow),
+  // otherwise returns RowVectorPtr with incremented values.
   RowVectorPtr createIncrementedBounds(const RowVectorPtr& bounds) const;
 
   // Encodes a single column for all rows.