Extract Column Logical Size from Raw Size Utils (#176)

Summary:

Rollback Plan:

Differential Revision: D75927903

Author: macvincent

dwio/nimble/velox/RawSizeContext.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include "dwio/nimble/common/Exceptions.h"
 #include "dwio/nimble/velox/DecodedVectorManager.h"
 
 namespace facebook::nimble {
@@ -28,8 +29,27 @@ class RawSizeContext {
     return decodedVectorManager_;
   }
 
+  void appendSize(uint64_t size) {
+    columnSizes_.push_back(size);
+  }
+
+  uint64_t sizeAt(uint64_t columnIndex) const {
+    NIMBLE_ASSERT(
+        columnIndex < columnSizes_.size(),
+        fmt::format(
+            "Column index {} is out of range. Total number of columns is {}",
+            columnIndex,
+            columnSizes_.size()));
+    return columnSizes_.at(columnIndex);
+  }
+
+  uint64_t columnCount() const {
+    return columnSizes_.size();
+  }
+
  private:
   DecodedVectorManager decodedVectorManager_;
+  std::vector<uint64_t> columnSizes_;
 };
 
 } // namespace facebook::nimble

dwio/nimble/velox/RawSizeUtils.cpp

@@ -442,7 +442,8 @@ uint64_t getRawSizeFromMapVector(
 uint64_t getRawSizeFromRowVector(
     const velox::VectorPtr& vector,
     const velox::common::Ranges& ranges,
-    RawSizeContext& context) {
+    RawSizeContext& context,
+    const bool topLevel) {
   VELOX_CHECK_NOT_NULL(vector);
   const auto& encoding = vector->encoding();
   const velox::RowVector* rowVector;
@@ -528,8 +529,12 @@ uint64_t getRawSizeFromRowVector(
   if ((*childRangesPtr).size()) {
     const auto childrenSize = rowVector->childrenSize();
     for (size_t i = 0; i < childrenSize; ++i) {
-      rawSize +=
+      auto childRawSize =
           getRawSizeFromVector(rowVector->childAt(i), *childRangesPtr, context);
+      rawSize += childRawSize;
+      if (topLevel) {
+        context.appendSize(childRawSize);
+      }
     }
   }
 

dwio/nimble/velox/RawSizeUtils.h

@@ -32,4 +32,9 @@ uint64_t getRawSizeFromVector(
     const velox::VectorPtr& vector,
     const velox::common::Ranges& ranges);
 
+uint64_t getRawSizeFromRowVector(
+    const velox::VectorPtr& vector,
+    const velox::common::Ranges& ranges,
+    RawSizeContext& context,
+    const bool topLevel = false);
 } // namespace facebook::nimble