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,21 @@ 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);
+    LOG(INFO) << "Raw size: " << rawSize;
     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 +598,26 @@ void VeloxWriter::close() {
              builder.GetSize()});
       }
 
+      {
+        // Accumulate column physical size.
+        std::vector<uint64_t> columnPhysicalSize(
+            context_->columnStats.size(), 0);
+        const auto& physicalSize = context_->streamPhysicalSize;
+        for (const auto& node : context_->schemaBuilder.getSchemaNodes()) {
+          auto column = node->column();
+          offset_size offset = node->offset();
+          if (physicalSize.contains(offset) && column.has_value()) {
+            NIMBLE_CHECK(
+                column.value() < columnPhysicalSize.size(),
+                fmt::format(
+                    "Column {} is out of range. Schema has {} columns.",
+                    column.value(),
+                    columnPhysicalSize.size()));
+            columnPhysicalSize[column.value()] += physicalSize.at(offset);
+          }
+        }
+      }
+
       {
         flatbuffers::FlatBufferBuilder builder;
         builder.Finish(serialization::CreateStats(builder, context_->rawSize));
@@ -691,14 +729,16 @@ void VeloxWriter::writeChunk(bool lastChunk) {
       StreamData& streamData_;
     };
 
-    auto encode = [&](StreamData& streamData) {
+    auto encode = [&](StreamData& streamData,
+                      std::atomic<uint64_t>& streamSize) {
       const auto offset = streamData.descriptor().offset();
       auto encoded = encodeStream(*context_, *encodingBuffer_, streamData);
       if (!encoded.empty()) {
         ChunkedStreamWriter chunkWriter{*encodingBuffer_};
         NIMBLE_DASSERT(offset < streams_.size(), "Stream offset out of range.");
         auto& stream = streams_[offset];
         for (auto& buffer : chunkWriter.encode(encoded)) {
+          streamSize += buffer.size();
           chunkSize += buffer.size();
           stream.content.push_back(std::move(buffer));
         }
@@ -739,29 +779,35 @@ void VeloxWriter::writeChunk(bool lastChunk) {
       velox::dwio::common::ExecutorBarrier barrier{
           context_->options.encodingExecutor};
       for (auto& streamData : context_->streams()) {
+        auto& streamSize =
+            context_->streamPhysicalSize[streamData->descriptor().offset()];
         processStream(
             *streamData, [&](StreamData& innerStreamData, bool isNullStream) {
-              barrier.add([&innerStreamData, isNullStream, &encode]() {
-                if (isNullStream) {
-                  NullsAsDataStreamData nullsStreamData{innerStreamData};
-                  encode(nullsStreamData);
-                } else {
-                  encode(innerStreamData);
-                }
-              });
+              barrier.add(
+                  [&innerStreamData, isNullStream, &encode, &streamSize]() {
+                    if (isNullStream) {
+                      NullsAsDataStreamData nullsStreamData{innerStreamData};
+                      encode(nullsStreamData, streamSize);
+                    } else {
+                      encode(innerStreamData, streamSize);
+                    }
+                  });
             });
       }
       barrier.waitAll();
     } else {
       for (auto& streamData : context_->streams()) {
+        auto& streamSize =
+            context_->streamPhysicalSize[streamData->descriptor().offset()];
         processStream(
             *streamData,
-            [&encode](StreamData& innerStreamData, bool isNullStream) {
+            [&encode, &streamSize](
+                StreamData& innerStreamData, bool isNullStream) {
               if (isNullStream) {
                 NullsAsDataStreamData nullsStreamData{innerStreamData};
-                encode(nullsStreamData);
+                encode(nullsStreamData, streamSize);
               } else {
-                encode(innerStreamData);
+                encode(innerStreamData, streamSize);
               }
             });
       }