Add a caller row group filter to RowReaderOptions

c++/include/orc/Reader.hh

@@ -28,6 +28,7 @@
 #include "orc/sargs/SearchArgument.hh"
 
 #include <atomic>
+#include <functional>
 #include <map>
 #include <memory>
 #include <set>
@@ -85,6 +86,17 @@ namespace orc {
     std::vector<std::vector<uint64_t>> positions;
   };
 
+  // Returns the statistics of one row group of one column, or nullptr when that column has no
+  // row index in the current stripe. The returned object is owned by the reader and is only
+  // valid for the duration of the RowGroupFilter call.
+  using RowGroupStatisticsAccessor =
+      std::function<const ColumnStatistics*(uint64_t columnId, uint64_t rowGroup)>;
+
+  // Caller-supplied row group selection, intersected with the reader's own search argument
+  // evaluation. See RowReaderOptions::rowGroupFilter.
+  using RowGroupFilter =
+      std::function<std::vector<bool>(uint64_t numRowGroups, const RowGroupStatisticsAccessor&)>;
+
   /**
    * Options for creating a Reader.
    */
@@ -283,6 +295,26 @@ namespace orc {
      */
     RowReaderOptions& searchArgument(std::shared_ptr<SearchArgument> sargs);
 
+    /**
+     * Set a callback that contributes an additional row group selection, evaluated by the
+     * caller against the row group statistics this reader has already loaded.
+     *
+     * The callback receives the number of row groups in the stripe and an accessor returning
+     * the statistics of a (column id, row group) pair, or nullptr when the column has no row
+     * index in this stripe. It returns one keep flag per row group; a vector of a different
+     * size is treated as all-keep.
+     *
+     * The result is intersected with this reader's own search argument evaluation, so the
+     * callback can only ever narrow the selection. Without a callback the reader behaves
+     * exactly as before.
+     */
+    RowReaderOptions& rowGroupFilter(RowGroupFilter filter);
+
+    /**
+     * Get the row group filter callback, empty if none was set.
+     */
+    const RowGroupFilter& getRowGroupFilter() const;
+
     /**
      * Should enable encoding block mode
      */

c++/src/Options.hh

@@ -150,6 +150,7 @@ namespace orc {
     int32_t forcedScaleOnHive11Decimal;
     bool enableLazyDecoding;
     std::shared_ptr<SearchArgument> sargs;
+    RowGroupFilter rowGroupFilter;
     std::string readerTimezone;
     RowReaderOptions::IdReadIntentMap idReadIntentMap;
     bool useTightNumericVector;
@@ -316,6 +317,15 @@ namespace orc {
     return privateBits_->sargs;
   }
 
+  RowReaderOptions& RowReaderOptions::rowGroupFilter(RowGroupFilter filter) {
+    privateBits_->rowGroupFilter = std::move(filter);
+    return *this;
+  }
+
+  const RowGroupFilter& RowReaderOptions::getRowGroupFilter() const {
+    return privateBits_->rowGroupFilter;
+  }
+
   RowReaderOptions& RowReaderOptions::setTimezoneName(const std::string& zoneName) {
     privateBits_->readerTimezone = zoneName;
     return *this;

c++/src/Reader.cc

@@ -369,6 +369,8 @@ namespace orc {
           *contents_->schema, sargs_.get(), footer_->row_index_stride(),
           getWriterVersionImpl(contents.get()), opts.getDictionaryFilteringSizeThreshold(),
           contents_->readerMetrics, &schemaEvolution_);
+      // The filter narrows the sargs selection, so it needs that applier to exist.
+      rowGroupFilter_ = opts.getRowGroupFilter();
     }
 
     skipBloomFilters_ = hasBadBloomFilters();
@@ -496,6 +498,38 @@ namespace orc {
     }
   }
 
+  std::vector<bool> RowReaderImpl::evaluateRowGroupFilter() {
+    uint64_t rowIndexStride = footer_->row_index_stride();
+    uint64_t numRowGroups = (rowsInCurrentStripe_ + rowIndexStride - 1) / rowIndexStride;
+
+    const Timezone& writerTimezone = currentStripeFooter_.has_writer_timezone()
+                                        ? getTimezoneByName(currentStripeFooter_.writer_timezone())
+                                        : getLocalTimezone();
+    StatContext statContext(
+        !WriterVersionImpl::VERSION_HIVE_8732().compareGT(getWriterVersionImpl(contents_.get())),
+        &writerTimezone);
+
+    // The wrappers are cached so that a filter reading one column across all row groups does
+    // not re-convert, and so that the pointers it receives stay valid for the whole call.
+    std::map<std::pair<uint64_t, uint64_t>, std::unique_ptr<ColumnStatistics>> converted;
+    auto accessor = [&](uint64_t columnId, uint64_t rowGroup) -> const ColumnStatistics* {
+      auto key = std::make_pair(columnId, rowGroup);
+      auto cached = converted.find(key);
+      if (cached != converted.end()) return cached->second.get();
+
+      auto rowIndex = rowIndexes_.find(columnId);
+      if (rowIndex == rowIndexes_.cend() ||
+          static_cast<int>(rowGroup) >= rowIndex->second.entry_size()) {
+        return nullptr;
+      }
+      std::unique_ptr<ColumnStatistics> stats(convertColumnStatistics(
+          rowIndex->second.entry(static_cast<int>(rowGroup)).statistics(), statContext));
+      return converted.emplace(key, std::move(stats)).first->second.get();
+    };
+
+    return rowGroupFilter_(numRowGroups, accessor);
+  }
+
   void RowReaderImpl::loadStripeIndex() {
     // reset all previous row indexes
     rowIndexes_.clear();
@@ -1295,8 +1329,15 @@ namespace orc {
           // read row group statistics and bloom filters of current stripe
           loadStripeIndex();
 
+          // ask the caller for its own row group selection, if it installed a filter
+          std::vector<bool> callerSelection;
+          if (rowGroupFilter_) {
+            callerSelection = evaluateRowGroupFilter();
+          }
+
           // select row groups to read in the current stripe
-          sargsApplier_->pickRowGroups(rowsInCurrentStripe_, rowIndexes_, bloomFilterIndex_);
+          sargsApplier_->pickRowGroups(rowsInCurrentStripe_, rowIndexes_, bloomFilterIndex_,
+                                       rowGroupFilter_ ? &callerSelection : nullptr);
           if (sargsApplier_->hasSelectedFrom(currentRowInStripe_)) {
             // current stripe has at least one row group matching the predicate
             break;

c++/src/Reader.hh

@@ -186,6 +186,8 @@ namespace orc {
     std::map<uint32_t, BloomFilterIndex> bloomFilterIndex_;
     std::shared_ptr<SearchArgument> sargs_;
     std::unique_ptr<SargsApplier> sargsApplier_;
+    // caller-supplied row group selection, intersected with the sargs evaluation
+    RowGroupFilter rowGroupFilter_;
 
     // desired timezone to return data of timestamp types.
     const Timezone& readerTimezone_;
@@ -199,6 +201,9 @@ namespace orc {
     // load stripe index if not done so
     void loadStripeIndex();
 
+    // run the caller's row group filter over the already-loaded row index of the current stripe
+    std::vector<bool> evaluateRowGroupFilter();
+
     // In case of PPD, batch size should be aware of row group boundaries.
     // If only a subset of row groups are selected then the next read should
     // stop at the end of selected range.

c++/src/sargs/SargsApplier.cc

@@ -20,6 +20,7 @@
 #include "Dictionary.hh"
 #include "sargs/PredicateLeaf.hh"
 
+#include <algorithm>
 #include <numeric>
 #include <set>
 
@@ -81,14 +82,28 @@ namespace orc {
 
   bool SargsApplier::pickRowGroups(uint64_t rowsInStripe,
                                    const std::unordered_map<uint64_t, proto::RowIndex>& rowIndexes,
-                                   const std::map<uint32_t, BloomFilterIndex>& bloomFilters) {
+                                   const std::map<uint32_t, BloomFilterIndex>& bloomFilters,
+                                   const std::vector<bool>* callerSelection) {
     // init state of each row group
     uint64_t groupsInStripe = (rowsInStripe + rowIndexStride_ - 1) / rowIndexStride_;
     nextSkippedRows_.resize(groupsInStripe);
     totalRowsInStripe_ = rowsInStripe;
 
+    // A selection of a different size cannot be matched to row groups; ignore it.
+    if (callerSelection != nullptr && callerSelection->size() != groupsInStripe) {
+      callerSelection = nullptr;
+    }
+
     // row indexes do not exist, simply read all rows
     if (rowIndexes.empty()) {
+      // With a caller selection installed the schedule must say "keep everything" explicitly:
+      // no statistics are available to evaluate it against, and an all-zero vector reads as
+      // "everything skipped" in hasSelectedFrom() and computeBatchSize().
+      if (callerSelection != nullptr) {
+        std::fill(nextSkippedRows_.begin(), nextSkippedRows_.end(), rowsInStripe);
+        hasSelected_ = groupsInStripe > 0;
+        hasSkipped_ = false;
+      }
       return true;
     }
 
@@ -125,7 +140,11 @@ namespace orc {
         }
       }
 
-      bool needed = isNeeded(searchArgument_->evaluate(leafValues));
+      // The caller's keep bit has to participate here, not in a later pass over
+      // nextSkippedRows_: that vector encodes the END of each selected run, so zeroing an
+      // entry afterwards would leave the preceding run's marker spanning the skipped group.
+      bool needed = isNeeded(searchArgument_->evaluate(leafValues)) &&
+                    (callerSelection == nullptr || (*callerSelection)[rowGroup]);
       if (!needed) {
         nextSkippedRows_[rowGroup] = 0;
         nextSkippedRowGroup = rowGroup;

c++/src/sargs/SargsApplier.hh

@@ -74,11 +74,15 @@ namespace orc {
     /**
      * TODO: use proto::RowIndex and proto::BloomFilter to do the evaluation
      * Pick the row groups that we need to load from the current stripe.
+     * @param callerSelection optional per-row-group keep flags contributed by the caller
+     *        (RowReaderOptions::rowGroupFilter), intersected with the search argument
+     *        evaluation. Ignored when null or of a size other than the row group count.
      * @return true if any row group is selected
      */
     bool pickRowGroups(uint64_t rowsInStripe,
                        const std::unordered_map<uint64_t, proto::RowIndex>& rowIndexes,
-                       const std::map<uint32_t, BloomFilterIndex>& bloomFilters);
+                       const std::map<uint32_t, BloomFilterIndex>& bloomFilters,
+                       const std::vector<bool>* callerSelection = nullptr);
 
     /**
      * Return a vector of the next skipped row for each RowGroup. Each value is the row id