add RowGroupPositions

apache · Sep 20, 2024 · f1b8e21 · f1b8e21
1 parent 7a31082
commit f1b8e21
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 13 deletions.
diff --git a/c++/include/orc/Reader.hh b/c++/include/orc/Reader.hh
@@ -62,6 +62,11 @@ namespace orc {
   };
   ReaderMetrics* getDefaultReaderMetrics();
 
+  struct RowGroupPositions {
+    uint64_t columnId;
+    std::vector<int32_t> positions;
+  };
+
   /**
    * Options for creating a Reader.
    */
@@ -660,9 +665,9 @@ namespace orc {
 
     /**
      * Get the row group positions of the specified column in the current stripe.
-     * @return the position entries for the specified column.
+     * @return the position entries for the specified columns.
      */
-    virtual std::vector<std::vector<int>> getCurrentStripePositionEntries(uint64_t columnId) = 0;
+    virtual std::vector<RowGroupPositions> getPositionEntries(int columnId) = 0;
   };
 }  // namespace orc
 

diff --git a/c++/src/Reader.cc b/c++/src/Reader.cc
@@ -1241,16 +1241,18 @@ namespace orc {
                                    useTightNumericVector_);
   }
 
-  std::vector<std::vector<int>> RowReaderImpl::getCurrentStripePositionEntries(uint64_t columnId) {
+  std::vector<RowGroupPositions> RowReaderImpl::getPositionEntries(int columnId) {
     loadStripeIndex();
-    std::vector<std::vector<int>> result;
+    std::vector<RowGroupPositions> result;
     auto rowIndex = rowIndexes_[columnId];
+    RowGroupPositions rgPositions;
+    rgPositions.columnId = columnId;
     for (auto rowIndexEntry : rowIndex.entry()) {
-      std::vector<int> posVector;
+      auto posVector = rgPositions.positions;
       for (auto position : rowIndexEntry.positions()) {
         posVector.push_back(position);
       }
-      result.push_back(posVector);
+      result.push_back(rgPositions);
     }
 
     return result;

diff --git a/c++/src/Reader.hh b/c++/src/Reader.hh
@@ -246,7 +246,7 @@ namespace orc {
       return &schemaEvolution_;
     }
 
-    std::vector<std::vector<int>> getCurrentStripePositionEntries(uint64_t columnId) override;
+    std::vector<RowGroupPositions> getPositionEntries(int columnId) override;
   };
 
   class ReaderImpl : public Reader {

diff --git a/c++/test/TestWriter.cc b/c++/test/TestWriter.cc
@@ -89,10 +89,10 @@ namespace orc {
   void verifyCompressionBlockAlignment(const std::unique_ptr<RowReader>& reader,
                                        uint64_t columnCount) {
     for (uint64_t i = 0; i < columnCount; ++i) {
-      auto posEntries = reader->getCurrentStripePositionEntries(i + 1);
+      auto posEntries = reader->getPositionEntries(i);
       auto subType = reader->getSelectedType().getSubtype(i);
-      for (auto posVector : posEntries) {
-        for (uint64_t posIndex = 0; posIndex < posVector.size(); ++posIndex) {
+      for (auto rowGroupPositions : posEntries) {
+        for (uint64_t posIndex = 0; posIndex < rowGroupPositions.positions.size(); ++posIndex) {
           // After we call finishStream(), unusedBufferSize is set to 0,
           // so only the first position is valid in each recordPosition call.
           switch (subType->getKind()) {
@@ -102,20 +102,20 @@ namespace orc {
             case CHAR:
             case VARCHAR: {
               if (posIndex != 0 && posIndex != 2) {
-                EXPECT_EQ(posVector[posIndex], 0);
+                EXPECT_EQ(rowGroupPositions.positions[posIndex], 0);
               }
               break;
             }
             case TIMESTAMP_INSTANT:
             case TIMESTAMP: {
               if (posIndex != 0 && posIndex != 3) {
-                EXPECT_EQ(posVector[posIndex], 0);
+                EXPECT_EQ(rowGroupPositions.positions[posIndex], 0);
               }
               break;
             }
             default: {
               if (posIndex != 0) {
-                EXPECT_EQ(posVector[posIndex], 0);
+                EXPECT_EQ(rowGroupPositions.positions[posIndex], 0);
               }
               break;
             }