Write Column Logical Size Into Nimble Files (#179)

Summary: Pull Request resolved: #179

Differential Revision: D75973860

Author: macvincent

dwio/nimble/velox/VeloxWriter.cpp

@@ -47,6 +47,11 @@ namespace detail {
 
 class WriterContext : public FieldWriterContext {
  public:
+  struct ColumnStats {
+    uint64_t logicalSize{0};
+    uint64_t nullCount{0};
+  };
+
   const VeloxWriterOptions options;
   std::unique_ptr<FlushPolicy> flushPolicy;
   velox::CpuWallTiming totalFlushTiming;
@@ -63,6 +68,8 @@ class WriterContext : public FieldWriterContext {
   uint64_t stripeSize{0};
   uint64_t rawSize{0};
   std::vector<uint64_t> rowsPerStripe;
+  std::map<offset_size, std::atomic<uint64_t>> streamPhysicalSize;
+  std::vector<ColumnStats> columnStats;
 
   WriterContext(
       velox::memory::MemoryPool& memoryPool,
@@ -516,10 +523,20 @@ 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);
+      context_->columnStats[i].nullCount += context.nullsAt(i);
+    }
 
     if (context_->options.writeExecutor) {
       velox::dwio::common::ExecutorBarrier barrier{
@@ -580,6 +597,32 @@ void VeloxWriter::close() {
              builder.GetSize()});
       }
 
+      {
+        // accumulate column physical size
+        // TODO: Clean up unnecessary logs
+        auto columnCount = context_->columnStats.size();
+        const auto& physicalSize = context_->streamPhysicalSize;
+        std::vector<uint64_t> columnPhysicalSize(columnCount, 0);
+        for (const auto& node : context_->schemaBuilder.getSchemaNodes()) {
+          auto column = node->parentNodeOffset();
+          offset_size offset = node->offset();
+          if (physicalSize.contains(offset) && column.has_value()) {
+            columnPhysicalSize[column.value()] += physicalSize.at(offset);
+          } else {
+            LOG(INFO) << "Nothing for stream " << offset;
+          }
+        }
+        // LOG column stats
+        for (int i = 0; i < columnCount; ++i) {
+          LOG(INFO) << fmt::format(
+              "Column {} logical size: {}, physical size: {}, null count: {}",
+              i,
+              context_->columnStats[i].logicalSize,
+              columnPhysicalSize[i],
+              context_->columnStats[i].nullCount);
+        }
+      }
+
       {
         flatbuffers::FlatBufferBuilder builder;
         builder.Finish(serialization::CreateStats(builder, context_->rawSize));
@@ -698,10 +741,13 @@ void VeloxWriter::writeChunk(bool lastChunk) {
         ChunkedStreamWriter chunkWriter{*encodingBuffer_};
         NIMBLE_DASSERT(offset < streams_.size(), "Stream offset out of range.");
         auto& stream = streams_[offset];
+        uint64_t streamSize = 0;
         for (auto& buffer : chunkWriter.encode(encoded)) {
-          chunkSize += buffer.size();
+          streamSize += buffer.size();
           stream.content.push_back(std::move(buffer));
         }
+        context_->streamPhysicalSize[offset].fetch_add(streamSize);
+        chunkSize += streamSize;
       }
       streamData.reset();
     };
@@ -778,8 +824,9 @@ void VeloxWriter::writeChunk(bool lastChunk) {
   auto flushWallTimeMs =
       (context_->stripeFlushTiming.wallNanos - previousFlushWallTime) /
       1'000'000;
-  VLOG(1) << "writeChunk milliseconds: " << flushWallTimeMs
-          << ", chunk bytes: " << chunkSize;
+  // TODO: Clean up unnecessary logs
+  LOG(INFO) << "writeChunk milliseconds: " << flushWallTimeMs
+            << ", chunk bytes: " << chunkSize;
 }
 
 uint32_t VeloxWriter::writeStripe() {