spillheap

Author: zuochunwei

velox/common/file/File.h

@@ -229,6 +229,164 @@ class InMemoryWriteFile final : public WriteFile {
   std::string* FOLLY_NONNULL file_;
 };
 
+// TODO zuochunwei
+struct HeapMemoryMock {
+  HeapMemoryMock() = default;
+  explicit HeapMemoryMock(void* memory, size_t capacity)
+      : memory_(memory), capacity_(capacity) {}
+
+  void reset() {
+    memory_ = nullptr;
+    size_ = 0;
+    capacity_ = 0;
+  }
+
+  bool isValid() const {
+    return memory_ != nullptr;
+  }
+
+  void write(const void* src, size_t len) {
+    assert(len <= freeSize());
+    memcpy(end(), src, len);
+    size_ += len;
+  }
+
+  void read(void* dst, size_t len, size_t offset) {
+    assert(offset + len <= size_);
+    memcpy(dst, (char*)memory_ + offset, len);
+  }
+
+  auto size() const {
+    return size_;
+  }
+
+  auto freeSize() const {
+    return capacity_ - size_;
+  }
+
+  void* begin() {
+    return memory_;
+  }
+
+  void* end() {
+    return (char*)memory_ + size_;
+  }
+
+  void* memory_ = nullptr;
+  size_t size_ = 0;
+  size_t capacity_ = 0;
+};
+
+const size_t kHeapMemoryCapacity = 64 * 1024;
+
+class HeapMemoryMockManager {
+ public:
+  static HeapMemoryMockManager& instance() {
+    static HeapMemoryMockManager hmmm;
+    return hmmm;
+  }
+
+  HeapMemoryMock alloc(size_t size) {
+    HeapMemoryMock heapMemory;
+    if (size_ + size <= kHeapMemoryCapacity) {
+      heapMemory.memory_ = malloc(size);
+      heapMemory.size_ = 0;
+      heapMemory.capacity_ = size;
+      size_ += size;
+    }
+    return heapMemory;
+  }
+
+  void free(HeapMemoryMock& heapMemory) {
+    if (heapMemory.isValid()) {
+      size_ -= heapMemory.size_;
+      ::free(heapMemory.memory_);
+      heapMemory.reset();
+    }
+  }
+
+ private:
+  std::atomic<std::size_t> size_;
+};
+
+inline HeapMemoryMock allocHeapMemory(size_t size) {
+  return HeapMemoryMockManager::instance().alloc(size);
+}
+
+inline void freeHeapMemory(HeapMemoryMock& heapMemory) {
+  HeapMemoryMockManager::instance().free(heapMemory);
+}
+
+class HeapMemoryReadFile : public ReadFile {
+ public:
+  explicit HeapMemoryReadFile(HeapMemoryMock& heapMemory)
+      : heapMemory_(heapMemory) {}
+
+  std::string_view pread(
+      uint64_t offset,
+      uint64_t length,
+      void* FOLLY_NONNULL buf) const override {
+    bytesRead_ += length;
+    heapMemory_.read(buf, length, offset);
+    return {static_cast<char*>(buf), length};
+  }
+
+  std::string pread(uint64_t offset, uint64_t length) const override {
+    bytesRead_ += length;
+    assert(offset + lenght <= heapMemory_.size());
+    return std::string((char*)heapMemory_.begin() + offset, length);
+  }
+
+  uint64_t size() const final {
+    return heapMemory_.size();
+  }
+
+  uint64_t memoryUsage() const final {
+    return size();
+  }
+
+  // Mainly for testing. Coalescing isn't helpful for in memory data.
+  void setShouldCoalesce(bool shouldCoalesce) {
+    shouldCoalesce_ = shouldCoalesce;
+  }
+  bool shouldCoalesce() const final {
+    return shouldCoalesce_;
+  }
+
+  std::string getName() const override {
+    return "<HeapMemoryReadFile>";
+  }
+
+  uint64_t getNaturalReadSize() const override {
+    return 1024;
+  }
+
+ private:
+  HeapMemoryMock& heapMemory_;
+  bool shouldCoalesce_ = false;
+};
+
+class HeapMemoryWriteFile final : public WriteFile {
+ public:
+  explicit HeapMemoryWriteFile(HeapMemoryMock& heapMemory)
+      : heapMemory_(heapMemory) {}
+
+  void append(std::string_view data) final {
+    heapMemory_.write(data.data(), data.length());
+  }
+
+  void flush() final {}
+
+  void close() final {}
+
+  uint64_t size() const final {
+    return heapMemory_.size_;
+  }
+
+ private:
+  HeapMemoryMock& heapMemory_;
+};
+
 // Current implementation for the local version is quite simple (e.g. no
 // internal arenaing), as local disk writes are expected to be cheap. Local
 // files match against any filepath starting with '/'.

velox/exec/Spill.cpp

@@ -45,24 +45,46 @@ void SpillMergeStream::pop() {
   }
 }
 
-WriteFile& SpillFile::output() {
-  if (!output_) {
+void SpillFile::newOutput() {
+  heapMemoryMock_ = allocHeapMemory(targetFileSize_);
+  if (heapMemoryMock_.isValid()) {
+    output_ = std::make_unique<HeapMemoryWriteFile>(heapMemoryMock_);
+    toWhere_ = TO_HEAP;
+  } else {
     auto fs = filesystems::getFileSystem(path_, nullptr);
     output_ = fs->openFileForWrite(path_);
+    toWhere_ = TO_FILE;
+  }
+}
+
+WriteFile& SpillFile::output() {
+  if (!output_) {
+    newOutput();
   }
   return *output_;
 }
 
 void SpillFile::startRead() {
   constexpr uint64_t kMaxReadBufferSize =
       (1 << 20) - AlignedBuffer::kPaddedSize; // 1MB - padding.
+
   VELOX_CHECK(!output_);
   VELOX_CHECK(!input_);
-  auto fs = filesystems::getFileSystem(path_, nullptr);
-  auto file = fs->openFileForRead(path_);
-  auto buffer = AlignedBuffer::allocate<char>(
-      std::min<uint64_t>(fileSize_, kMaxReadBufferSize), &pool_);
-  input_ = std::make_unique<SpillInput>(std::move(file), std::move(buffer));
+
+  if (toWhere_ == TO_FILE) {
+    auto fs = filesystems::getFileSystem(path_, nullptr);
+    auto file = fs->openFileForRead(path_);
+    auto buffer = AlignedBuffer::allocate<char>(
+        std::min<uint64_t>(fileSize_, kMaxReadBufferSize), &pool_);
+    input_ = std::make_unique<SpillInput>(std::move(file), std::move(buffer));
+  } else if (toWhere_ == TO_HEAP) {
+    auto file = std::make_unique<HeapMemoryReadFile>(heapMemoryMock_);
+    auto buffer = AlignedBuffer::allocate<char>(
+        std::min<uint64_t>(fileSize_, kMaxReadBufferSize), &pool_);
+    input_ = std::make_unique<SpillInput>(std::move(file), std::move(buffer));
+  } else {
+    VELOX_FAIL("invalid spill destination");
+  }
 }
 
 bool SpillFile::nextBatch(RowVectorPtr& rowVector) {
@@ -74,18 +96,20 @@ bool SpillFile::nextBatch(RowVectorPtr& rowVector) {
   return true;
 }
 
-WriteFile& SpillFileList::currentOutput() {
+WriteFile& SpillFileList::currentOutput(size_t toAppendSize) {
   if (files_.empty() || !files_.back()->isWritable() ||
-      files_.back()->size() > targetFileSize_) {
+      files_.back()->size() + toAppendSize > targetFileSize_) {
     if (!files_.empty() && files_.back()->isWritable()) {
       files_.back()->finishWrite();
     }
+    assert(toAppendSize <= targetFileSize_);
     files_.push_back(std::make_unique<SpillFile>(
         type_,
         numSortingKeys_,
         sortCompareFlags_,
         fmt::format("{}-{}", path_, files_.size()),
-        pool_));
+        pool_,
+        targetFileSize_));
   }
   return files_.back()->output();
 }
@@ -97,7 +121,14 @@ void SpillFileList::flush() {
     batch_->flush(&out);
     batch_.reset();
     auto iobuf = out.getIOBuf();
-    auto& file = currentOutput();
+
+    size_t toAppendSize = 0;
+    for (auto& range : *iobuf) {
+      toAppendSize += range.size();
+    }
+
+    auto& file = currentOutput(toAppendSize);
+
     for (auto& range : *iobuf) {
       file.append(std::string_view(
           reinterpret_cast<const char*>(range.data()), range.size()));

velox/exec/Spill.h

@@ -68,20 +68,28 @@ class SpillFile {
       int32_t numSortingKeys,
       const std::vector<CompareFlags>& sortCompareFlags,
       const std::string& path,
-      memory::MemoryPool& pool)
+      memory::MemoryPool& pool,
+      uint64_t targetFileSize)
       : type_(std::move(type)),
         numSortingKeys_(numSortingKeys),
         sortCompareFlags_(sortCompareFlags),
         pool_(pool),
         ordinal_(ordinalCounter_++),
-        path_(fmt::format("{}-{}", path, ordinal_)) {
+        path_(fmt::format("{}-{}", path, ordinal_)),
+        targetFileSize_(targetFileSize) {
     // NOTE: if the spilling operator has specified the sort comparison flags,
     // then it must match the number of sorting keys.
     VELOX_CHECK(
         sortCompareFlags_.empty() ||
         sortCompareFlags_.size() == numSortingKeys_);
   }
 
+  ~SpillFile() {
+    if (heapMemoryMock_.isValid()) {
+      freeHeapMemory(heapMemoryMock_);
+    }
+  }
+
   int32_t numSortingKeys() const {
     return numSortingKeys_;
   }
@@ -133,6 +141,8 @@ class SpillFile {
   }
 
  private:
+  void newOutput();
+
   static std::atomic<int32_t> ordinalCounter_;
 
   // Type of 'rowVector_'. Needed for setting up writing.
@@ -145,8 +155,17 @@ class SpillFile {
   const int32_t ordinal_;
   const std::string path_;
 
+  enum {
+    TO_FILE,
+    TO_HEAP,
+  } toWhere_ = TO_FILE;
+
+  HeapMemoryMock heapMemoryMock_;
+
   // Byte size of the backing file. Set when finishing writing.
   uint64_t fileSize_ = 0;
+  uint64_t targetFileSize_ = 0;
+
   std::unique_ptr<WriteFile> output_;
   std::unique_ptr<SpillInput> input_;
 };
@@ -215,7 +234,7 @@ class SpillFileList {
 
  private:
   // Returns the current file to write to and creates one if needed.
-  WriteFile& currentOutput();
+  WriteFile& currentOutput(size_t toAppendSize);
 
   // Writes data from 'batch_' to the current output file.
   void flush();