Write Column Logical Size Into Nimble Files

Differential Revision: D75973860

Author: macvincent

dwio/nimble/velox/VeloxWriter.cpp

@@ -47,6 +47,10 @@ namespace detail {
 
 class WriterContext : public FieldWriterContext {
  public:
+  struct ColumnStats {
+    uint64_t logicalSize{0};
+  };
+
   const VeloxWriterOptions options;
   std::unique_ptr<FlushPolicy> flushPolicy;
   velox::CpuWallTiming totalFlushTiming;
@@ -63,6 +67,7 @@ class WriterContext : public FieldWriterContext {
   uint64_t stripeSize{0};
   uint64_t rawSize{0};
   std::vector<uint64_t> rowsPerStripe;
+  std::vector<ColumnStats> columnStats;
 
   WriterContext(
       velox::memory::MemoryPool& memoryPool,
@@ -516,10 +521,19 @@ bool VeloxWriter::write(const velox::VectorPtr& vector) {
     auto size = vector->size();
 
     // Calculate raw size.
-    auto rawSize = nimble::getRawSizeFromVector(
-        vector, velox::common::Ranges::of(0, size));
+    RawSizeContext context;
+    auto rawSize = nimble::getRawSizeFromRowVector(
+        vector, velox::common::Ranges::of(0, size), context, /*topLevel=*/true);
     DWIO_ENSURE_GE(rawSize, 0, "Invalid raw size");
     context_->rawSize += rawSize;
+    auto columnCount = context.columnCount();
+    if (context_->columnStats.empty()) {
+      context_->columnStats =
+          std::vector<detail::WriterContext::ColumnStats>(columnCount);
+    }
+    for (auto i = 0; i < columnCount; ++i) {
+      context_->columnStats[i].logicalSize += context.sizeAt(i);
+    }
 
     if (context_->options.writeExecutor) {
       velox::dwio::common::ExecutorBarrier barrier{