Add physicalWrittenBytes into the PlanNodeStats (facebookincubator#8213)

Pass spark 3.4 unit test when enabling native parquet write (#466)

velox/connectors/hive/HiveConfig.cpp

@@ -153,6 +153,11 @@ bool HiveConfig::isFileColumnNamesReadAsLowerCase(const Config* session) const {
   return config_->get<bool>(kFileColumnNamesReadAsLowerCase, false);
 }
 
+bool HiveConfig::isPartitionPathAsLowerCaseSession(
+    const Config* session) const {
+  return config_->get<bool>(kPartitionPathAsLowerCaseSession, true);
+}
+
 int64_t HiveConfig::maxCoalescedBytes() const {
   return config_->get<int64_t>(kMaxCoalescedBytes, 128 << 20);
 }

velox/connectors/hive/HiveConfig.h

@@ -110,6 +110,9 @@ class HiveConfig {
   static constexpr const char* kFileColumnNamesReadAsLowerCaseSession =
       "file_column_names_read_as_lower_case";
 
+  static constexpr const char* kPartitionPathAsLowerCaseSession =
+      "partition_path_as_lower_case";
+
   /// Sets the max coalesce bytes for a request.
   static constexpr const char* kMaxCoalescedBytes = "max-coalesced-bytes";
 
@@ -202,6 +205,8 @@ class HiveConfig {
 
   bool isFileColumnNamesReadAsLowerCase(const Config* session) const;
 
+  bool isPartitionPathAsLowerCaseSession(const Config* session) const;
+
   int64_t maxCoalescedBytes() const;
 
   int32_t maxCoalescedDistanceBytes() const;

velox/connectors/hive/HiveDataSink.cpp

@@ -38,6 +38,28 @@ namespace facebook::velox::connector::hive {
 
 namespace {
 
+RowVectorPtr makeDataInput(
+    const std::vector<column_index_t>& partitonCols,
+    const RowVectorPtr& input,
+    const RowTypePtr& dataType) {
+  std::vector<VectorPtr> childVectors;
+  childVectors.reserve(dataType->size());
+  for (uint32_t i = 0; i < input->childrenSize(); i++) {
+    if (std::find(partitonCols.cbegin(), partitonCols.cend(), i) ==
+        partitonCols.cend()) {
+      childVectors.push_back(input->childAt(i));
+    }
+  }
+
+  return std::make_shared<RowVector>(
+      input->pool(),
+      dataType,
+      input->nulls(),
+      input->size(),
+      std::move(childVectors),
+      input->getNullCount());
+}
+
 // Returns a subset of column indices corresponding to partition keys.
 std::vector<column_index_t> getPartitionChannels(
     const std::shared_ptr<const HiveInsertTableHandle>& insertTableHandle) {
@@ -301,12 +323,15 @@ HiveDataSink::HiveDataSink(
           connectorQueryCtx->sessionProperties())),
       partitionChannels_(getPartitionChannels(insertTableHandle_)),
       partitionIdGenerator_(
-          !partitionChannels_.empty() ? std::make_unique<PartitionIdGenerator>(
-                                            inputType_,
-                                            partitionChannels_,
-                                            maxOpenWriters_,
-                                            connectorQueryCtx_->memoryPool())
-                                      : nullptr),
+          !partitionChannels_.empty()
+              ? std::make_unique<PartitionIdGenerator>(
+                    inputType_,
+                    partitionChannels_,
+                    maxOpenWriters_,
+                    connectorQueryCtx_->memoryPool(),
+                    hiveConfig_->isPartitionPathAsLowerCaseSession(
+                        connectorQueryCtx->sessionProperties()))
+              : nullptr),
       bucketCount_(
           insertTableHandle_->bucketProperty() == nullptr
               ? 0
@@ -331,6 +356,18 @@ HiveDataSink::HiveDataSink(
       "Unsupported commit strategy: {}",
       commitStrategyToString(commitStrategy_));
 
+  // Get the data input type based on the inputType and the parition index.
+  std::vector<TypePtr> childTypes;
+  std::vector<std::string> childNames;
+  for (auto i = 0; i < inputType_->size(); i++) {
+    if (std::find(partitionChannels_.cbegin(), partitionChannels_.cend(), i) ==
+        partitionChannels_.end()) {
+      childNames.push_back(inputType_->nameOf(i));
+      childTypes.push_back(inputType_->childAt(i));
+    }
+  }
+  dataType_ = ROW(std::move(childNames), std::move(childTypes));
+
   if (!isBucketed()) {
     return;
   }
@@ -400,8 +437,17 @@ void HiveDataSink::appendData(RowVectorPtr input) {
 
 void HiveDataSink::write(size_t index, const VectorPtr& input) {
   WRITER_NON_RECLAIMABLE_SECTION_GUARD(index);
-  writers_[index]->write(input);
-  writerInfo_[index]->numWrittenRows += input->size();
+  // Skip the partition columns before writing.
+  auto dataInput = input;
+  if (!isBucketed()) {
+    dataInput = makeDataInput(
+        partitionChannels_,
+        std::dynamic_pointer_cast<RowVector>(input),
+        dataType_);
+  }
+
+  writers_[index]->write(dataInput);
+  writerInfo_[index]->numWrittenRows += dataInput->size();
 }
 
 std::string HiveDataSink::stateString(State state) {
@@ -597,7 +643,12 @@ uint32_t HiveDataSink::appendWriter(const HiveWriterId& id) {
   dwio::common::WriterOptions options;
   const auto* connectorSessionProperties =
       connectorQueryCtx_->sessionProperties();
-  options.schema = inputType_;
+  if (!isBucketed()) {
+    options.schema = dataType_;
+  } else {
+    options.schema = inputType_;
+  }
+
   options.memoryPool = writerInfo_.back()->writerPool.get();
   options.compressionKind = insertTableHandle_->compressionKind();
   if (canReclaim()) {

velox/connectors/hive/HiveDataSink.h

@@ -541,6 +541,8 @@ class HiveDataSink : public DataSink {
   void closeInternal();
 
   const RowTypePtr inputType_;
+  // Written data columns into file.
+  RowTypePtr dataType_;
   const std::shared_ptr<const HiveInsertTableHandle> insertTableHandle_;
   const ConnectorQueryCtx* const connectorQueryCtx_;
   const CommitStrategy commitStrategy_;

velox/connectors/hive/PartitionIdGenerator.cpp

@@ -27,9 +27,11 @@ PartitionIdGenerator::PartitionIdGenerator(
     const RowTypePtr& inputType,
     std::vector<column_index_t> partitionChannels,
     uint32_t maxPartitions,
-    memory::MemoryPool* pool)
+    memory::MemoryPool* pool,
+    bool partitionPathAsLowerCase)
     : partitionChannels_(std::move(partitionChannels)),
-      maxPartitions_(maxPartitions) {
+      maxPartitions_(maxPartitions),
+      partitionPathAsLowerCase_(partitionPathAsLowerCase) {
   VELOX_USER_CHECK(
       !partitionChannels_.empty(), "There must be at least one partition key.");
   for (auto channel : partitionChannels_) {
@@ -98,7 +100,8 @@ void PartitionIdGenerator::run(
 
 std::string PartitionIdGenerator::partitionName(uint64_t partitionId) const {
   return FileUtils::makePartName(
-      extractPartitionKeyValues(partitionValues_, partitionId));
+      extractPartitionKeyValues(partitionValues_, partitionId),
+      partitionPathAsLowerCase_);
 }
 
 void PartitionIdGenerator::computeValueIds(

velox/connectors/hive/PartitionIdGenerator.h

@@ -34,7 +34,8 @@ class PartitionIdGenerator {
       const RowTypePtr& inputType,
       std::vector<column_index_t> partitionChannels,
       uint32_t maxPartitions,
-      memory::MemoryPool* pool);
+      memory::MemoryPool* pool,
+      bool partitionPathAsLowerCase = true);
 
   /// Generate sequential partition IDs for input vector.
   /// @param input Input RowVector.
@@ -77,6 +78,8 @@ class PartitionIdGenerator {
 
   const uint32_t maxPartitions_;
 
+  bool partitionPathAsLowerCase_;
+
   std::vector<std::unique_ptr<exec::VectorHasher>> hashers_;
 
   // A mapping from value ID produced by VectorHashers to a partition ID.

velox/dwio/catalog/fbhive/FileUtils.cpp

@@ -158,14 +158,20 @@ std::string FileUtils::unescapePathName(const std::string& data) {
 }
 
 std::string FileUtils::makePartName(
-    const std::vector<std::pair<std::string, std::string>>& entries) {
+    const std::vector<std::pair<std::string, std::string>>& entries,
+    bool partitionPathAsLowerCase) {
   size_t size = 0;
   size_t escapeCount = 0;
   std::for_each(entries.begin(), entries.end(), [&](auto& pair) {
     auto keySize = pair.first.size();
     DWIO_ENSURE_GT(keySize, 0);
     size += keySize;
-    escapeCount += countEscape(pair.first);
+    if (partitionPathAsLowerCase) {
+      escapeCount += countEscape(toLower(pair.first));
+    } else {
+      escapeCount += countEscape(pair.first);
+    }
+
     auto valSize = pair.second.size();
     if (valSize == 0) {
       size += DEFAULT_PARTITION_VALUE.size();
@@ -182,7 +188,7 @@ std::string FileUtils::makePartName(
     if (ret.size() > 0) {
       ret += "/";
     }
-    ret += escapePathName(toLower(pair.first));
+    ret += escapePathName(pair.first);
     ret += "=";
     if (pair.second.size() == 0) {
       ret += DEFAULT_PARTITION_VALUE;

velox/dwio/catalog/fbhive/FileUtils.h

@@ -40,7 +40,8 @@ class FileUtils {
   /// Creates the partition directory path from the list of partition key/value
   /// pairs, will do url-encoding when needed.
   static std::string makePartName(
-      const std::vector<std::pair<std::string, std::string>>& entries);
+      const std::vector<std::pair<std::string, std::string>>& entries,
+      bool partitionPathAsLowerCase = true);
 
   /// Converts the hive-metastore-compliant path name back to the corresponding
   /// partition key/value pairs.

velox/dwio/catalog/fbhive/test/FileUtilsTests.cpp

@@ -28,7 +28,7 @@ TEST(FileUtilsTests, MakePartName) {
       {"ds", "2016-01-01"}, {"FOO", ""}, {"a\nb:c", "a#b=c"}};
   ASSERT_EQ(
       FileUtils::makePartName(pairs),
-      "ds=2016-01-01/foo=__HIVE_DEFAULT_PARTITION__/a%0Ab%3Ac=a%23b%3Dc");
+      "ds=2016-01-01/FOO=__HIVE_DEFAULT_PARTITION__/a%0Ab%3Ac=a%23b%3Dc");
 }
 
 TEST(FileUtilsTests, ParsePartKeyValues) {

velox/exec/PlanNodeStats.cpp

@@ -54,6 +54,8 @@ void PlanNodeStats::addTotals(const OperatorStats& stats) {
   peakMemoryBytes += stats.memoryStats.peakTotalMemoryReservation;
   numMemoryAllocations += stats.memoryStats.numMemoryAllocations;
 
+  physicalWrittenBytes += stats.physicalWrittenBytes;
+
   for (const auto& [name, runtimeStats] : stats.runtimeStats) {
     if (UNLIKELY(customStats.count(name) == 0)) {
       customStats.insert(std::make_pair(name, runtimeStats));
@@ -154,6 +156,7 @@ folly::dynamic toPlanStatsJson(const facebook::velox::exec::TaskStats& stats) {
       stat["blockedWallNanos"] = operatorStat.second->blockedWallNanos;
       stat["peakMemoryBytes"] = operatorStat.second->peakMemoryBytes;
       stat["numMemoryAllocations"] = operatorStat.second->numMemoryAllocations;
+      stat["physicalWrittenBytes"] = operatorStat.second->physicalWrittenBytes;
       stat["numDrivers"] = operatorStat.second->numDrivers;
       stat["numSplits"] = operatorStat.second->numSplits;
       stat["spilledInputBytes"] = operatorStat.second->spilledInputBytes;

velox/exec/PlanNodeStats.h

@@ -92,6 +92,8 @@ struct PlanNodeStats {
 
   uint64_t numMemoryAllocations{0};
 
+  uint64_t physicalWrittenBytes{0};
+
   /// Operator-specific counters.
   std::unordered_map<std::string, RuntimeMetric> customStats;