diff --git a/c++/include/orc/Reader.hh b/c++/include/orc/Reader.hh index 81cd583e2c..c11c264b52 100644 --- a/c++/include/orc/Reader.hh +++ b/c++/include/orc/Reader.hh @@ -28,6 +28,7 @@ #include "orc/sargs/SearchArgument.hh" #include +#include #include #include #include @@ -85,6 +86,17 @@ namespace orc { std::vector> 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; + + // Caller-supplied row group selection, intersected with the reader's own search argument + // evaluation. See RowReaderOptions::rowGroupFilter. + using RowGroupFilter = + std::function(uint64_t numRowGroups, const RowGroupStatisticsAccessor&)>; + /** * Options for creating a Reader. */ @@ -283,6 +295,26 @@ namespace orc { */ RowReaderOptions& searchArgument(std::shared_ptr 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 */ diff --git a/c++/src/Options.hh b/c++/src/Options.hh index daa46336ed..e6a6abb012 100644 --- a/c++/src/Options.hh +++ b/c++/src/Options.hh @@ -150,6 +150,7 @@ namespace orc { int32_t forcedScaleOnHive11Decimal; bool enableLazyDecoding; std::shared_ptr 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; diff --git a/c++/src/Reader.cc b/c++/src/Reader.cc index d5054fdf99..c0780f861f 100644 --- a/c++/src/Reader.cc +++ b/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 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::unique_ptr> 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(rowGroup) >= rowIndex->second.entry_size()) { + return nullptr; + } + std::unique_ptr stats(convertColumnStatistics( + rowIndex->second.entry(static_cast(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 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; diff --git a/c++/src/Reader.hh b/c++/src/Reader.hh index 204f678e96..2d8d98503f 100644 --- a/c++/src/Reader.hh +++ b/c++/src/Reader.hh @@ -186,6 +186,8 @@ namespace orc { std::map bloomFilterIndex_; std::shared_ptr sargs_; std::unique_ptr 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 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. diff --git a/c++/src/sargs/SargsApplier.cc b/c++/src/sargs/SargsApplier.cc index 5c7aa10efc..2091a2e10e 100644 --- a/c++/src/sargs/SargsApplier.cc +++ b/c++/src/sargs/SargsApplier.cc @@ -20,6 +20,7 @@ #include "Dictionary.hh" #include "sargs/PredicateLeaf.hh" +#include #include #include @@ -81,14 +82,28 @@ namespace orc { bool SargsApplier::pickRowGroups(uint64_t rowsInStripe, const std::unordered_map& rowIndexes, - const std::map& bloomFilters) { + const std::map& bloomFilters, + const std::vector* 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; diff --git a/c++/src/sargs/SargsApplier.hh b/c++/src/sargs/SargsApplier.hh index b9e85d5080..ffedd56212 100644 --- a/c++/src/sargs/SargsApplier.hh +++ b/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& rowIndexes, - const std::map& bloomFilters); + const std::map& bloomFilters, + const std::vector* callerSelection = nullptr); /** * Return a vector of the next skipped row for each RowGroup. Each value is the row id