fix: Iceberg commit message when writer rotated (facebookincubator#16365)

Summary:
Fixes commit message correctness for multi-file writes by updating Iceberg writer rotation handling and tracking per-file row counts and emitting one commit task per rotated file.
Fix parquet writer does not honor smaller values of `max-target-file-size` config.

These are the adjustments required after facebookincubator#16077.

Pull Request resolved: facebookincubator#16365

Reviewed By: Yuhta

Differential Revision: D93157234

Pulled By: jainxrohit

fbshipit-source-id: f09cd1a01cb2ee63172f233d9fc8234688037fb6

Author: PingLiuPing

velox/connectors/hive/HiveDataSink.cpp

@@ -621,6 +621,7 @@ void HiveDataSink::write(size_t index, RowVectorPtr input) {
   writers_[index]->write(dataInput);
   writerInfo_[index]->inputSizeInBytes += dataInput->estimateFlatSize();
   writerInfo_[index]->numWrittenRows += dataInput->size();
+  writerInfo_[index]->currentFileWrittenRows += dataInput->size();
 
   // File rotation is not supported for bucketed tables (require one file per
   // bucket with predictable name) or sorted writes (SortingWriter not
@@ -642,12 +643,7 @@ uint64_t HiveDataSink::getCurrentFileBytes(size_t writerIndex) const {
   const auto baselineBytes = writerInfo_[writerIndex]->cumulativeWrittenBytes;
   // Sanity check: total should always be >= baseline since ioStats is
   // never reset and cumulative is a snapshot of rawBytesWritten at rotation.
-  VELOX_DCHECK_GE(
-      totalBytes,
-      baselineBytes,
-      "rawBytesWritten ({}) < cumulativeWrittenBytes ({})",
-      totalBytes,
-      baselineBytes);
+  VELOX_DCHECK_GE(totalBytes, baselineBytes);
   return totalBytes - baselineBytes;
 }
 
@@ -665,6 +661,9 @@ void HiveDataSink::finalizeWriterFile(size_t index) {
   fileInfo.writeFileName = info->currentWriteFileName;
   fileInfo.targetFileName = info->currentTargetFileName;
   fileInfo.fileSize = currentFileBytes;
+  fileInfo.numRows = info->currentFileWrittenRows;
+  // Reset for next file.
+  info->currentFileWrittenRows = 0;
   info->writtenFiles.push_back(std::move(fileInfo));
 
   // Update cumulative stats as a snapshot of total stats so far.

velox/connectors/hive/HiveDataSink.h

@@ -426,6 +426,8 @@ struct HiveFileInfo {
   std::string targetFileName;
   /// Size of the file in bytes.
   uint64_t fileSize{0};
+  /// Number of rows in the file.
+  uint64_t numRows{0};
 };
 
 struct HiveWriterInfo {
@@ -449,8 +451,11 @@ struct HiveWriterInfo {
   const std::shared_ptr<memory::MemoryPool> writerPool;
   const std::shared_ptr<memory::MemoryPool> sinkPool;
   const std::shared_ptr<memory::MemoryPool> sortPool;
-  int64_t numWrittenRows = 0;
-  int64_t inputSizeInBytes = 0;
+  /// Total rows written by this writer across all files.
+  uint64_t numWrittenRows = 0;
+  /// Rows written to the current file; reset to 0 when the file is finalized.
+  uint64_t currentFileWrittenRows{0};
+  uint64_t inputSizeInBytes = 0;
   /// File sequence number for tracking multiple files written due to size-based
   /// splitting. Incremented each time the writer rotates to a new file.
   /// Used to generate sequenced file names (e.g., file_1.orc, file_2.orc).

velox/connectors/hive/iceberg/IcebergDataSink.cpp

@@ -315,30 +315,36 @@ std::vector<std::string> IcebergDataSink::commitMessage() const {
           insertTableHandle_);
 
   for (auto i = 0; i < writerInfo_.size(); ++i) {
-    const auto& info = writerInfo_.at(i);
-    VELOX_CHECK_NOT_NULL(info);
+    const auto& writerInfo = writerInfo_.at(i);
+    VELOX_CHECK_NOT_NULL(writerInfo);
+
     // Following metadata (json format) is consumed by Presto CommitTaskData.
     // It contains the minimal subset of metadata.
     // TODO: Complete metrics is missing now and this could lead to suboptimal
     // query plan, will collect full iceberg metrics in following PR.
-    // clang-format off
-    folly::dynamic commitData = folly::dynamic::object(
-    "path", (fs::path(info->writerParameters.writeDirectory()) /
-                    info->writerParameters.writeFileName()).string())
-      ("fileSizeInBytes", ioStats_.at(i)->rawBytesWritten())
-      ("metrics",
-        folly::dynamic::object("recordCount", info->numWrittenRows))
-      ("partitionSpecJson",
-        icebergInsertTableHandle->partitionSpec() ? icebergInsertTableHandle->partitionSpec()->specId : 0)
-      ("fileFormat", "PARQUET")
-      ("content", "DATA");
-    // clang-format on
-    if (!commitPartitionValue_.empty() && !commitPartitionValue_[i].isNull()) {
-      commitData["partitionDataJson"] = folly::toJson(
-          folly::dynamic::object("partitionValues", commitPartitionValue_[i]));
+    for (const auto& fileInfo : writerInfo->writtenFiles) {
+      // clang-format off
+      folly::dynamic commitData = folly::dynamic::object(
+        "path", (fs::path(writerInfo->writerParameters.targetDirectory()) /
+                      fileInfo.targetFileName).string())
+        ("fileSizeInBytes", fileInfo.fileSize)
+        ("metrics",
+          folly::dynamic::object("recordCount", fileInfo.numRows))
+        ("partitionSpecJson",
+          icebergInsertTableHandle->partitionSpec() ?
+            icebergInsertTableHandle->partitionSpec()->specId : 0)
+        ("fileFormat", "PARQUET")
+        ("content", "DATA");
+      // clang-format on
+      if (!commitPartitionValue_.empty() &&
+          !commitPartitionValue_[i].isNull()) {
+        commitData["partitionDataJson"] = folly::toJson(
+            folly::dynamic::object(
+                "partitionValues", commitPartitionValue_[i]));
+      }
+      auto commitDataJson = folly::toJson(commitData);
+      commitTasks.push_back(commitDataJson);
     }
-    auto commitDataJson = folly::toJson(commitData);
-    commitTasks.push_back(commitDataJson);
   }
   return commitTasks;
 }

velox/connectors/hive/iceberg/tests/TransformE2ETest.cpp

@@ -280,14 +280,14 @@ class TransformE2ETest : public test::IcebergTestBase {
     const auto dataSink = createDataSinkAndAppendData(
         {rowVector}, outputDirectory->getPath(), partitionTransforms);
 
+    dataSink->close();
     auto commitMessages = dataSink->commitMessage();
     VELOX_CHECK_EQ(commitMessages.size(), 1);
     auto commitData = folly::parseJson(commitMessages[0]);
     auto partitionDataJson =
         folly::parseJson(commitData["partitionDataJson"].asString());
     auto partitionValues = partitionDataJson["partitionValues"];
     VELOX_CHECK_EQ(partitionValues.size(), 1);
-    dataSink->close();
     return partitionValues[0];
   }
 };