ORC-1264: [C++] Add a writer option to align compression block with row group

c++/include/orc/Reader.hh

@@ -62,9 +62,13 @@ namespace orc {
   };
   ReaderMetrics* getDefaultReaderMetrics();
 
-  struct RowGroupPositions {
-    uint64_t columnId;
-    std::vector<int32_t> positions;
+  // Row group index of a single column in a stripe.
+  struct RowGroupIndex {
+    // Positions are represented as a two-dimensional array where the first
+    // dimension is row group index and the second dimension is the position
+    // list of the row group. The size of the second dimension should be equal
+    // among all row groups.
+    std::vector<std::vector<uint64_t>> positions;
   };
 
   /**
@@ -610,6 +614,16 @@ namespace orc {
      */
     virtual std::map<uint32_t, BloomFilterIndex> getBloomFilters(
         uint32_t stripeIndex, const std::set<uint32_t>& included) const = 0;
+
+    /**
+     * Get row group index of all selected columns in the specified stripe
+     * @param stripeIndex index of the stripe to be read for row group index.
+     * @param included index of selected columns to return (if not specified,
+     *        all columns will be returned).
+     * @return map of row group index keyed by its column index.
+     */
+    virtual std::map<uint32_t, RowGroupIndex> getRowGroupIndex(
+        uint32_t stripeIndex, const std::set<uint32_t>& included = {}) const = 0;
   };
 
   /**
@@ -663,11 +677,6 @@ namespace orc {
      */
     virtual void seekToRow(uint64_t rowNumber) = 0;
 
-    /**
-     * Get the row group positions of the specified column in the current stripe.
-     * @return the position entries for the specified columns.
-     */
-    virtual std::vector<RowGroupPositions> getPositionEntries(int columnId) = 0;
   };
 }  // namespace orc
 

c++/src/Reader.cc

@@ -1241,23 +1241,6 @@ namespace orc {
                                    useTightNumericVector_);
   }
 
-  std::vector<RowGroupPositions> RowReaderImpl::getPositionEntries(int columnId) {
-    loadStripeIndex();
-    std::vector<RowGroupPositions> result;
-    auto rowIndex = rowIndexes_[columnId];
-    RowGroupPositions rgPositions;
-    rgPositions.columnId = columnId;
-    for (auto rowIndexEntry : rowIndex.entry()) {
-      auto posVector = rgPositions.positions;
-      for (auto position : rowIndexEntry.positions()) {
-        posVector.push_back(position);
-      }
-      result.push_back(rgPositions);
-    }
-
-    return result;
-  }
-
   void ensureOrcFooter(InputStream* stream, DataBuffer<char>* buffer, uint64_t postscriptLength) {
     const std::string MAGIC("ORC");
     const uint64_t magicLength = MAGIC.length();
@@ -1443,17 +1426,10 @@ namespace orc {
       uint32_t stripeIndex, const std::set<uint32_t>& included) const {
     std::map<uint32_t, BloomFilterIndex> ret;
 
-    // find stripe info
-    if (stripeIndex >= static_cast<uint32_t>(footer_->stripes_size())) {
-      throw std::logic_error("Illegal stripe index: " +
-                             to_string(static_cast<int64_t>(stripeIndex)));
-    }
-    const proto::StripeInformation currentStripeInfo =
-        footer_->stripes(static_cast<int>(stripeIndex));
-    const proto::StripeFooter currentStripeFooter = getStripeFooter(currentStripeInfo, *contents_);
+    uint64_t offset;
+    auto currentStripeFooter = loadCurrentStripeFooter(stripeIndex, offset);
 
     // iterate stripe footer to get stream of bloom_filter
-    uint64_t offset = static_cast<uint64_t>(currentStripeInfo.offset());
     for (int i = 0; i < currentStripeFooter.streams_size(); i++) {
       const proto::Stream& stream = currentStripeFooter.streams(i);
       uint32_t column = static_cast<uint32_t>(stream.column());
@@ -1491,6 +1467,59 @@ namespace orc {
     return ret;
   }
 
+  proto::StripeFooter ReaderImpl::loadCurrentStripeFooter(uint32_t stripeIndex,
+                                                          uint64_t& offset) const {
+    // find stripe info
+    if (stripeIndex >= static_cast<uint32_t>(footer_->stripes_size())) {
+      throw std::logic_error("Illegal stripe index: " +
+                             to_string(static_cast<int64_t>(stripeIndex)));
+    }
+    const proto::StripeInformation currentStripeInfo =
+        footer_->stripes(static_cast<int>(stripeIndex));
+    offset = static_cast<uint64_t>(currentStripeInfo.offset());
+    return getStripeFooter(currentStripeInfo, *contents_);
+  }
+
+  std::map<uint32_t, RowGroupIndex> ReaderImpl::getRowGroupIndex(
+      uint32_t stripeIndex, const std::set<uint32_t>& included) const {
+    std::map<uint32_t, RowGroupIndex> ret;
+    uint64_t offset;
+    auto currentStripeFooter = loadCurrentStripeFooter(stripeIndex, offset);
+
+    // iterate stripe footer to get stream of row_index
+    for (int i = 0; i < currentStripeFooter.streams_size(); i++) {
+      const proto::Stream& stream = currentStripeFooter.streams(i);
+      uint32_t column = static_cast<uint32_t>(stream.column());
+      uint64_t length = static_cast<uint64_t>(stream.length());
+      RowGroupIndex rowGroupIndex = ret[column];
+
+      if (stream.kind() == proto::Stream_Kind_ROW_INDEX &&
+          (included.empty() || included.find(column) != included.end())) {
+        std::unique_ptr<SeekableInputStream> pbStream =
+            createDecompressor(contents_->compression,
+                               std::make_unique<SeekableFileInputStream>(
+                                   contents_->stream.get(), offset, length, *contents_->pool),
+                               contents_->blockSize, *(contents_->pool), contents_->readerMetrics);
+
+        proto::RowIndex pbRowIndex;
+        if (!pbRowIndex.ParseFromZeroCopyStream(pbStream.get())) {
+          throw ParseError("Failed to parse RowIndex");
+        }
+
+        // add rowGroupIndex to result for one column
+        for (auto& rowIndexEntry : pbRowIndex.entry()) {
+          std::vector<uint64_t> posVector;
+          for (auto& position : rowIndexEntry.positions()) {
+            posVector.push_back(position);
+          }
+          rowGroupIndex.positions.push_back(posVector);
+        }
+      }
+      offset += length;
+    }
+    return ret;
+  }
+
   RowReader::~RowReader() {
     // PASS
   }
@@ -1499,8 +1528,8 @@ namespace orc {
     // PASS
   }
 
-  InputStream::~InputStream(){
-      // PASS
+  InputStream::~InputStream() {
+    // PASS
   };
 
 }  // namespace orc

c++/src/Reader.hh

@@ -245,8 +245,6 @@ namespace orc {
     const SchemaEvolution* getSchemaEvolution() const {
       return &schemaEvolution_;
     }
-
-    std::vector<RowGroupPositions> getPositionEntries(int columnId) override;
   };
 
   class ReaderImpl : public Reader {
@@ -270,6 +268,7 @@ namespace orc {
     void getRowIndexStatistics(const proto::StripeInformation& stripeInfo, uint64_t stripeIndex,
                                const proto::StripeFooter& currentStripeFooter,
                                std::vector<std::vector<proto::ColumnStatistics>>* indexStats) const;
+    proto::StripeFooter loadCurrentStripeFooter(uint32_t stripeIndex, uint64_t& offset) const;
 
     // metadata
     mutable bool isMetadataLoaded_;
@@ -375,6 +374,9 @@ namespace orc {
 
     std::map<uint32_t, BloomFilterIndex> getBloomFilters(
         uint32_t stripeIndex, const std::set<uint32_t>& included) const override;
+
+    std::map<uint32_t, RowGroupIndex> getRowGroupIndex(
+        uint32_t stripeIndex, const std::set<uint32_t>& included) const override;
   };
 }  // namespace orc