Serialization Mix

velox/connectors/hive/HiveConnector.cpp

@@ -60,7 +60,7 @@ HiveConnector::HiveConnector(
                     SimpleLRUCache<std::string, std::shared_ptr<FileHandle>>>(
                     hiveConfig_->numCacheFileHandles())
               : nullptr,
-          std::make_unique<FileHandleGenerator>(config)),
+          std::make_unique<FileHandleGenerator>(nullptr)),
       executor_(executor) {
   if (hiveConfig_->isFileHandleCacheEnabled()) {
     LOG(INFO) << "Hive connector " << connectorId()

velox/connectors/hive/storage_adapters/s3fs/tests/CMakeLists.txt

@@ -64,20 +64,3 @@ target_link_libraries(
   velox_exec
   gtest
   gtest_main)
-
-add_executable(velox_s3read_test S3ReadTest.cpp)
-add_test(
-  NAME velox_s3read_test
-  COMMAND velox_s3read_test
-  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR})
-target_link_libraries(
-  velox_s3read_test
-  velox_file
-  velox_s3fs
-  velox_hive_config
-  velox_core
-  velox_exec_test_lib
-  velox_dwio_common_exception
-  velox_exec
-  gtest
-  gtest_main)

velox/connectors/hive/storage_adapters/s3fs/tests/MinioServer.h

@@ -14,8 +14,6 @@
  * limitations under the License.
  */
 
-#pragma once
-
 #include "velox/core/Config.h"
 #include "velox/exec/tests/utils/TempDirectoryPath.h"
 

velox/core/QueryConfig.h

@@ -156,9 +156,6 @@ class QueryConfig {
   static constexpr const char* kMaxPartitionedOutputBufferSize =
       "max_page_partitioning_buffer_size";
 
-  static constexpr const char* kMaxArbitraryBufferSize =
-      "max_arbitrary_buffer_size";
-
   /// Preferred size of batches in bytes to be returned by operators from
   /// Operator::getOutput. It is used when an estimate of average row size is
   /// known. Otherwise kPreferredOutputBatchRows is used.
@@ -397,26 +394,15 @@ class QueryConfig {
     return get<uint64_t>(kOrderBySpillMemoryThreshold, kDefault);
   }
 
-  /// Returns the maximum size in bytes for the task's buffered output when
-  /// output is partitioned using hash of partitioning keys. See
-  /// PartitionedOutputNode::Kind::kPartitioned.
-  ///
-  /// The producer Drivers are blocked when the buffered size exceeds
-  /// this. The Drivers are resumed when the buffered size goes below
-  /// OutputBufferManager::kContinuePct % of this.
+  // Returns the target size for a Task's buffered output. The
+  // producer Drivers are blocked when the buffered size exceeds
+  // this. The Drivers are resumed when the buffered size goes below
+  // OutputBufferManager::kContinuePct % of this.
   uint64_t maxPartitionedOutputBufferSize() const {
     static constexpr uint64_t kDefault = 32UL << 20;
     return get<uint64_t>(kMaxPartitionedOutputBufferSize, kDefault);
   }
 
-  /// Returns the maximum size in bytes for the task's buffered output when
-  /// output is distributed randomly among consumers. See
-  /// PartitionedOutputNode::Kind::kArbitrary.
-  uint64_t maxArbitraryBufferSize() const {
-    static constexpr uint64_t kDefault = 32UL << 20;
-    return get<uint64_t>(kMaxArbitraryBufferSize, kDefault);
-  }
-
   uint64_t maxLocalExchangeBufferSize() const {
     static constexpr uint64_t kDefault = 32UL << 20;
     return get<uint64_t>(kMaxLocalExchangeBufferSize, kDefault);

velox/docs/configs.rst

@@ -94,14 +94,7 @@ Generic Configuration
    * - max_page_partitioning_buffer_size
      - integer
      - 32MB
-     - The maximum size in bytes for the task's buffered output when output is partitioned using hash of partitioning keys. See PartitionedOutputNode::Kind::kPartitioned.
-       The producer Drivers are blocked when the buffered size exceeds this.
-       The Drivers are resumed when the buffered size goes below OutputBufferManager::kContinuePct (90)% of this.
-   * - max_arbitrary_buffer_size
-     - integer
-     - 32MB
-     - The maximum size in bytes for the task's buffered output when output is distributed randomly among consumers. See PartitionedOutputNode::Kind::kArbitrary.
-       The producer Drivers are blocked when the buffered size exceeds this.
+     - The target size for a Task's buffered output. The producer Drivers are blocked when the buffered size exceeds this.
        The Drivers are resumed when the buffered size goes below OutputBufferManager::kContinuePct (90)% of this.
    * - min_table_rows_for_parallel_join_build
      - integer

velox/exec/OutputBuffer.cpp

@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 #include "velox/exec/OutputBuffer.h"
-#include "velox/core/QueryConfig.h"
 #include "velox/exec/Task.h"
 
 namespace facebook::velox::exec {
@@ -218,17 +217,6 @@ void releaseAfterAcknowledge(
     promise.setValue();
   }
 }
-
-uint64_t maxBufferSize(
-    const core::QueryConfig& config,
-    PartitionedOutputNode::Kind bufferKind) {
-  if (bufferKind == PartitionedOutputNode::Kind::kArbitrary) {
-    return config.maxArbitraryBufferSize();
-  }
-
-  return config.maxPartitionedOutputBufferSize();
-}
-
 } // namespace
 
 OutputBuffer::OutputBuffer(
@@ -238,7 +226,8 @@ OutputBuffer::OutputBuffer(
     uint32_t numDrivers)
     : task_(std::move(task)),
       kind_(kind),
-      maxSize_(maxBufferSize(task_->queryCtx()->queryConfig(), kind)),
+      maxSize_(
+          task_->queryCtx()->queryConfig().maxPartitionedOutputBufferSize()),
       continueSize_((maxSize_ * kContinuePct) / 100),
       arbitraryBuffer_(
           isArbitrary() ? std::make_unique<ArbitraryBuffer>() : nullptr),
@@ -306,9 +295,9 @@ void OutputBuffer::addOutputBuffersLocked(int numBuffers) {
       for (const auto& data : dataToBroadcast_) {
         buffer->enqueue(data);
       }
-      if (atEnd_) {
-        buffer->enqueue(nullptr);
-      }
+    }
+    if (atEnd_ && isBroadcast()) {
+      buffer->enqueue(nullptr);
     }
     buffers_.emplace_back(std::move(buffer));
   }

velox/exec/PartitionedOutput.cpp

@@ -28,36 +28,27 @@ BlockingReason Destination::advance(
     OutputBufferManager& bufferManager,
     const std::function<void()>& bufferReleaseFn,
     bool* atEnd,
-    ContinueFuture* future) {
-  if (rangeIdx_ >= ranges_.size()) {
+    ContinueFuture* future,
+    Scratch& scratch) {
+  if (rowIdx_ >= rows_.size()) {
     *atEnd = true;
     return BlockingReason::kNotBlocked;
   }
 
+  const auto firstRow = rowIdx_;
   const uint32_t adjustedMaxBytes = (maxBytes * targetSizePct_) / 100;
   if (bytesInCurrent_ >= adjustedMaxBytes) {
     return flush(bufferManager, bufferReleaseFn, future);
   }
 
-  // Collect ranges to serialize.
-  rangesToSerialize_.clear();
+  // Collect rows to serialize.
   bool shouldFlush = false;
-  while (rangeIdx_ < ranges_.size() && !shouldFlush) {
-    auto& currRange = ranges_[rangeIdx_];
-    auto startRow = rowsInCurrentRange_;
-    for (; rowsInCurrentRange_ < currRange.size && !shouldFlush;
-         rowsInCurrentRange_++) {
-      ++rowsInCurrent_;
-      bytesInCurrent_ += sizes[currRange.begin + rowsInCurrentRange_];
-      shouldFlush = bytesInCurrent_ >= adjustedMaxBytes ||
-          rowsInCurrent_ >= targetNumRows_;
-    }
-    rangesToSerialize_.push_back(
-        {currRange.begin + startRow, rowsInCurrentRange_ - startRow});
-    if (rowsInCurrentRange_ == currRange.size) {
-      rowsInCurrentRange_ = 0;
-      rangeIdx_++;
-    }
+  while (rowIdx_ < rows_.size() && !shouldFlush) {
+    bytesInCurrent_ += sizes[rowIdx_];
+    ++rowIdx_;
+    ++rowsInCurrent_;
+    shouldFlush =
+        bytesInCurrent_ >= adjustedMaxBytes || rowsInCurrent_ >= targetNumRows_;
   }
 
   // Serialize
@@ -67,9 +58,9 @@ BlockingReason Destination::advance(
     current_->createStreamTree(rowType, rowsInCurrent_);
   }
   current_->append(
-      output, folly::Range(&rangesToSerialize_[0], rangesToSerialize_.size()));
+      output, folly::Range(&rows_[firstRow], rowIdx_ - firstRow), scratch);
   // Update output state variable.
-  if (rangeIdx_ == ranges_.size()) {
+  if (rowIdx_ == rows_.size()) {
     *atEnd = true;
   }
   if (shouldFlush || (eagerFlush_ && rowsInCurrent_ > 0)) {
@@ -189,12 +180,7 @@ void PartitionedOutput::initializeDestinations() {
 
 void PartitionedOutput::initializeSizeBuffers() {
   auto numInput = input_->size();
-  if (numInput > topLevelRanges_.size()) {
-    vector_size_t numOld = topLevelRanges_.size();
-    topLevelRanges_.resize(numInput);
-    for (auto i = numOld; i < numInput; ++i) {
-      topLevelRanges_[i] = IndexRange{i, 1};
-    }
+  if (numInput > rowSize_.size()) {
     rowSize_.resize(numInput);
     sizePointers_.resize(numInput);
     // Set all the size pointers since 'rowSize_' may have been reallocated.
@@ -207,11 +193,14 @@ void PartitionedOutput::initializeSizeBuffers() {
 void PartitionedOutput::estimateRowSizes() {
   auto numInput = input_->size();
   std::fill(rowSize_.begin(), rowSize_.end(), 0);
+  raw_vector<vector_size_t> storage;
+  auto numbers = iota(numInput, storage);
   for (int i = 0; i < output_->childrenSize(); ++i) {
     VectorStreamGroup::estimateSerializedSize(
         output_->childAt(i),
-        folly::Range(topLevelRanges_.data(), numInput),
-        sizePointers_.data());
+        folly::Range(numbers, numInput),
+        sizePointers_.data(),
+        scratch_);
   }
 }
 
@@ -332,7 +321,8 @@ RowVectorPtr PartitionedOutput::getOutput() {
           *bufferManager,
           bufferReleaseFn_,
           &atEnd,
-          &future_);
+          &future_,
+          scratch_);
       if (blockingReason_ != BlockingReason::kNotBlocked) {
         blockedDestination = destination.get();
         workLeft = false;