feat: Resubmit add max scan rows limit for index lookup (facebookincubator#16400)

Summary:
Pull Request resolved: facebookincubator#16400

X-link: facebookincubator/nimble#492

Original commit changeset: 5846b9ddaf6c

Original Phabricator Diff: D93313029

Reviewed By: HuamengJiang

Differential Revision: D93348953

fbshipit-source-id: a70b687917429542b3564993834cb1e61518eb7f

Author: xiaoxmeng

velox/connectors/hive/HiveConfig.cpp

@@ -259,6 +259,13 @@ bool HiveConfig::preserveFlatMapsInMemory(
       config_->get<bool>(kPreserveFlatMapsInMemory, false));
 }
 
+uint32_t HiveConfig::maxRowsPerIndexRequest(
+    const config::ConfigBase* session) const {
+  return session->get<uint32_t>(
+      kMaxRowsPerIndexRequestSession,
+      config_->get<uint32_t>(kMaxRowsPerIndexRequest, 0));
+}
+
 std::string HiveConfig::user(const config::ConfigBase* session) const {
   return session->get<std::string>(kUser, config_->get<std::string>(kUser, ""));
 }

velox/connectors/hive/HiveConfig.h

@@ -208,6 +208,14 @@ class HiveConfig {
   static constexpr const char* kPreserveFlatMapsInMemorySession =
       "hive.preserve_flat_maps_in_memory";
 
+  /// Maximum number of output rows to return per index lookup request.
+  /// The limit is applied to the actual output rows after filtering.
+  /// 0 means no limit (default).
+  static constexpr const char* kMaxRowsPerIndexRequest =
+      "hive.max-rows-per-index-request";
+  static constexpr const char* kMaxRowsPerIndexRequestSession =
+      "hive.max_rows_per_index_request";
+
   static constexpr const char* kUser = "user";
   static constexpr const char* kSource = "source";
   static constexpr const char* kSchema = "schema";
@@ -293,6 +301,10 @@ class HiveConfig {
   /// converting them to MapVectors.
   bool preserveFlatMapsInMemory(const config::ConfigBase* session) const;
 
+  /// Returns the maximum number of rows to read per index lookup request.
+  /// 0 means no limit (default).
+  uint32_t maxRowsPerIndexRequest(const config::ConfigBase* session) const;
+
   /// User of the query. Used for storage logging.
   std::string user(const config::ConfigBase* session) const;
 

velox/connectors/hive/HiveIndexReader.cpp

@@ -302,7 +302,9 @@ HiveIndexReader::createIndexReader() {
   return fileReader_->createIndexReader(rowReaderOpts);
 }
 
-void HiveIndexReader::startLookup(const Request& request) {
+void HiveIndexReader::startLookup(
+    const Request& request,
+    const Options& options) {
   VELOX_CHECK(
       !indexReader_->hasNext(),
       "Previous request not finished. Call next() first.");
@@ -311,7 +313,7 @@ void HiveIndexReader::startLookup(const Request& request) {
 
   // Build index bounds from request and pass to the index reader.
   auto indexBounds = buildRequestIndexBounds(request.input);
-  indexReader_->startLookup(indexBounds);
+  indexReader_->startLookup(indexBounds, options);
 }
 
 serializer::IndexBounds HiveIndexReader::buildRequestIndexBounds(

velox/connectors/hive/HiveIndexReader.h

@@ -63,13 +63,16 @@ class HiveIndexReader {
 
   using Request = IndexSource::Request;
   using Result = IndexSource::Result;
+  using Options = dwio::common::IndexReader::Options;
 
   /// Sets the input for index lookup. Each row in 'input' will be converted
   /// to index bounds and passed to the format-specific IndexReader.
   ///
   /// @param request The lookup request containing input row vector with lookup
   /// keys.
-  void startLookup(const Request& request);
+  /// @param options Options controlling index reader behavior (e.g.,
+  ///        maxRowsPerRequest). Defaults to no limit.
+  void startLookup(const Request& request, const Options& options = {});
 
   /// Returns true if there are more results to fetch from the current lookup.
   bool hasNext() const;

velox/connectors/hive/HiveIndexSource.cpp

@@ -227,10 +227,12 @@ class HiveLookupIterator : public IndexSource::ResultIterator {
   HiveLookupIterator(
       std::shared_ptr<HiveIndexSource> indexSource,
       HiveIndexReader* indexReader,
-      IndexSource::Request request)
+      IndexSource::Request request,
+      vector_size_t maxRowsPerRequest)
       : indexSource_(std::move(indexSource)),
         indexReader_(indexReader),
-        request_(std::move(request)) {}
+        request_(std::move(request)),
+        maxRowsPerRequest_(maxRowsPerRequest) {}
 
   ~HiveLookupIterator() override = default;
 
@@ -247,7 +249,8 @@ class HiveLookupIterator : public IndexSource::ResultIterator {
 
     // Set the request on first call.
     if (state_ == State::kInit) {
-      indexReader_->startLookup(request_);
+      indexReader_->startLookup(
+          request_, {.maxRowsPerRequest = maxRowsPerRequest_});
       setState(State::kRead);
     }
 
@@ -360,11 +363,11 @@ class HiveLookupIterator : public IndexSource::ResultIterator {
         emptyResult_->inputHits, emptyResult_->output);
   }
 
-  // Holds the index source to ensure HiveIndexReader lifetime.
   const std::shared_ptr<HiveIndexSource> indexSource_;
   // Raw pointer to index reader for lookup operations.
   HiveIndexReader* const indexReader_;
   const IndexSource::Request request_;
+  const vector_size_t maxRowsPerRequest_;
 
   State state_{State::kInit};
   // Cached empty result for reuse when no rows pass the remaining filter.
@@ -386,6 +389,8 @@ HiveIndexSource::HiveIndexSource(
       hiveConfig_(hiveConfig),
       pool_(connectorQueryCtx->memoryPool()),
       expressionEvaluator_(connectorQueryCtx->expressionEvaluator()),
+      maxRowsPerIndexRequest_(hiveConfig_->maxRowsPerIndexRequest(
+          connectorQueryCtx_->sessionProperties())),
       tableHandle_(std::move(tableHandle)),
       requestType_(requestType),
       outputType_(outputType),
@@ -653,7 +658,7 @@ std::shared_ptr<IndexSource::ResultIterator> HiveIndexSource::lookup(
     const Request& request) {
   VELOX_CHECK_NOT_NULL(indexReader_, "No index reader available for lookup");
   return std::make_shared<HiveLookupIterator>(
-      shared_from_this(), indexReader_.get(), request);
+      shared_from_this(), indexReader_.get(), request, maxRowsPerIndexRequest_);
 }
 
 std::unordered_map<std::string, RuntimeMetric> HiveIndexSource::runtimeStats() {

velox/connectors/hive/HiveIndexSource.h

@@ -68,6 +68,10 @@ class HiveIndexSource : public IndexSource,
     return pool_;
   }
 
+  uint32_t maxRowsPerIndexRequest() const {
+    return maxRowsPerIndexRequest_;
+  }
+
   const RowTypePtr& outputType() const {
     return outputType_;
   }
@@ -121,6 +125,7 @@ class HiveIndexSource : public IndexSource,
   const std::shared_ptr<HiveConfig> hiveConfig_;
   memory::MemoryPool* const pool_;
   core::ExpressionEvaluator* const expressionEvaluator_;
+  const uint32_t maxRowsPerIndexRequest_;
 
   const HiveTableHandlePtr tableHandle_;
   const RowTypePtr requestType_;

velox/docs/configs.rst

@@ -791,6 +791,12 @@ Each query can override the config by setting corresponding query session proper
      - bool
      - false
      - Whether to preserve flat maps in memory as FlatMapVectors instead of converting them to MapVectors. This is only applied during data reading inside the DWRF and Nimble readers, not during downstream processing like expression evaluation etc.
+   * - hive.max-rows-per-index-request
+     - hive.max_rows_per_index_request
+     - integer
+     - 0
+     - Maximum number of output rows to return per index lookup request. The limit is applied to the actual output rows
+       after filtering. 0 means no limit (default).
 
 ``ORC File Format Configuration``
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

velox/dwio/common/Reader.h

@@ -229,16 +229,28 @@ class IndexReader {
 
   virtual ~IndexReader() = default;
 
+  /// Options for controlling index reader behavior.
+  struct Options {
+    /// Maximum number of rows to read per index lookup request.
+    /// When set to non-zero, the index reader will stop fetching or truncate
+    /// stripes once the total row range (before filtering) reaches this limit.
+    /// 0 means no limit (default).
+    vector_size_t maxRowsPerRequest{0};
+  };
+
   /// Starts a new batch lookup with the given index bounds.
   /// Each index bound in the vector represents a separate lookup request.
   /// After calling startLookup(), call next() repeatedly to get results.
   ///
   /// @param indexBounds Index bounds for the lookup request. Contains
   ///        column names and lower/upper bound values.
+  /// @param options Options controlling index reader behavior (e.g.,
+  ///        maxRowsPerRequest). Defaults to no limit.
   /// @throws if lookup is not supported by the implementation or if any
   ///         index bound is invalid.
   virtual void startLookup(
-      const velox::serializer::IndexBounds& indexBounds) = 0;
+      const velox::serializer::IndexBounds& indexBounds,
+      const Options& options) = 0;
 
   /// Returns true if there are more results to fetch from the current lookup.
   virtual bool hasNext() const = 0;