Optimize PagedInputStream::Skip (facebookincubator#6699)

Summary:
Pull Request resolved: facebookincubator#6699

Currently when we skip bytes in `PagedInputStream`, we do the decompression unconditionally and it is expensive.  Some optimizations are added to address this:
1. Skip decompression of the whole block (frame in case of ZSTD) if
   1. We can get the precise decompressed size, and
   2. The decompressed size is no larger than the bytes need to skip
2. Accumulate contiguous skip calls to create larger skip region (delayed skipping)

Reviewed By: oerling

Differential Revision: D49501856

fbshipit-source-id: 07241aaf71e83f0f491050a9be6075dd5500dd52

Author: Yuhta

velox/dwio/common/compression/Compression.cpp

@@ -246,24 +246,32 @@ class ZstdDecompressor : public Decompressor {
   explicit ZstdDecompressor(
       uint64_t blockSize,
       const std::string& streamDebugInfo)
-      : Decompressor{blockSize, streamDebugInfo} {}
+      : Decompressor{blockSize, streamDebugInfo}, context_{ZSTD_createDCtx()} {}
+
+  ~ZstdDecompressor() override {
+    ZSTD_freeDCtx(context_);
+  }
 
   uint64_t decompress(
       const char* src,
       uint64_t srcLength,
       char* dest,
       uint64_t destLength) override;
 
-  uint64_t getUncompressedLength(const char* src, uint64_t srcLength)
-      const override;
+  std::pair<int64_t, bool> getDecompressedLength(
+      const char* src,
+      uint64_t srcLength) const override;
+
+ private:
+  ZSTD_DCtx* context_;
 };
 
 uint64_t ZstdDecompressor::decompress(
     const char* src,
     uint64_t srcLength,
     char* dest,
     uint64_t destLength) {
-  auto ret = ZSTD_decompress(dest, destLength, src, srcLength);
+  auto ret = ZSTD_decompressDCtx(context_, dest, destLength, src, srcLength);
   DWIO_ENSURE(
       !ZSTD_isError(ret),
       "ZSTD returned an error: ",
@@ -273,22 +281,22 @@ uint64_t ZstdDecompressor::decompress(
   return ret;
 }
 
-uint64_t ZstdDecompressor::getUncompressedLength(
+std::pair<int64_t, bool> ZstdDecompressor::getDecompressedLength(
     const char* src,
     uint64_t srcLength) const {
   auto uncompressedLength = ZSTD_getFrameContentSize(src, srcLength);
   // in the case when decompression size is not available, return the upper
   // bound
   if (uncompressedLength == ZSTD_CONTENTSIZE_UNKNOWN ||
       uncompressedLength == ZSTD_CONTENTSIZE_ERROR) {
-    return blockSize_;
+    return {blockSize_, false};
   }
   DWIO_ENSURE_LE(
       uncompressedLength,
       blockSize_,
       "Insufficient buffer size. Info: ",
       streamDebugInfo_);
-  return uncompressedLength;
+  return {uncompressedLength, true};
 }
 
 class SnappyDecompressor : public Decompressor {
@@ -304,16 +312,17 @@ class SnappyDecompressor : public Decompressor {
       char* dest,
       uint64_t destLength) override;
 
-  uint64_t getUncompressedLength(const char* src, uint64_t srcLength)
-      const override;
+  std::pair<int64_t, bool> getDecompressedLength(
+      const char* src,
+      uint64_t srcLength) const override;
 };
 
 uint64_t SnappyDecompressor::decompress(
     const char* src,
     uint64_t srcLength,
     char* dest,
     uint64_t destLength) {
-  auto length = getUncompressedLength(src, srcLength);
+  auto [length, _] = getDecompressedLength(src, srcLength);
   DWIO_ENSURE_GE(destLength, length);
   DWIO_ENSURE(
       snappy::RawUncompress(src, srcLength, dest),
@@ -322,23 +331,24 @@ uint64_t SnappyDecompressor::decompress(
   return length;
 }
 
-uint64_t SnappyDecompressor::getUncompressedLength(
+std::pair<int64_t, bool> SnappyDecompressor::getDecompressedLength(
     const char* src,
     uint64_t srcLength) const {
   size_t uncompressedLength;
   // in the case when decompression size is not available, return the upper
   // bound
   if (!snappy::GetUncompressedLength(src, srcLength, &uncompressedLength)) {
-    return blockSize_;
+    return {blockSize_, false};
   }
   DWIO_ENSURE_LE(
       uncompressedLength,
       blockSize_,
       "Insufficient buffer size. Info: ",
       streamDebugInfo_);
-  return uncompressedLength;
+  return {uncompressedLength, true};
 }
 
+// TODO: Is this really needed?
 class ZlibDecompressionStream : public PagedInputStream,
                                 private ZlibDecompressor {
  public:
@@ -355,13 +365,18 @@ class ZlibDecompressionStream : public PagedInputStream,
         ZlibDecompressor{blockSize, windowBits, streamDebugInfo, isGzip} {}
   ~ZlibDecompressionStream() override = default;
 
-  bool Next(const void** data, int32_t* size) override;
+  bool readOrSkip(const void** data, int32_t* size) override;
 };
 
-bool ZlibDecompressionStream::Next(const void** data, int32_t* size) {
+bool ZlibDecompressionStream::readOrSkip(const void** data, int32_t* size) {
+  if (data) {
+    VELOX_CHECK_EQ(pendingSkip_, 0);
+  }
   // if the user pushed back, return them the partial buffer
   if (outputBufferLength_) {
-    *data = outputBufferPtr_;
+    if (data) {
+      *data = outputBufferPtr_;
+    }
     *size = static_cast<int32_t>(outputBufferLength_);
     outputBufferPtr_ += outputBufferLength_;
     bytesReturned_ += outputBufferLength_;
@@ -381,7 +396,9 @@ bool ZlibDecompressionStream::Next(const void** data, int32_t* size) {
       static_cast<size_t>(inputBufferPtrEnd_ - inputBufferPtr_),
       remainingLength_);
   if (state_ == State::ORIGINAL) {
-    *data = inputBufferPtr_;
+    if (data) {
+      *data = inputBufferPtr_;
+    }
     *size = static_cast<int32_t>(availSize);
     outputBufferPtr_ = inputBufferPtr_ + availSize;
     outputBufferLength_ = 0;
@@ -393,7 +410,8 @@ bool ZlibDecompressionStream::Next(const void** data, int32_t* size) {
         getName(),
         " Info: ",
         ZlibDecompressor::streamDebugInfo_);
-    prepareOutputBuffer(getUncompressedLength(inputBufferPtr_, availSize));
+    prepareOutputBuffer(
+        getDecompressedLength(inputBufferPtr_, availSize).first);
 
     reset();
     zstream_.next_in =
@@ -432,7 +450,9 @@ bool ZlibDecompressionStream::Next(const void** data, int32_t* size) {
       }
     } while (result != Z_STREAM_END);
     *size = static_cast<int32_t>(blockSize_ - zstream_.avail_out);
-    *data = outputBufferPtr_;
+    if (data) {
+      *data = outputBufferPtr_;
+    }
     outputBufferLength_ = 0;
     outputBufferPtr_ += *size;
   }

velox/dwio/common/compression/Compression.h

@@ -46,14 +46,14 @@ class Compressor {
 class Decompressor {
  public:
   explicit Decompressor(uint64_t blockSize, const std::string& streamDebugInfo)
-      : blockSize_{blockSize}, streamDebugInfo_{streamDebugInfo} {}
+      : blockSize_{static_cast<int64_t>(blockSize)},
+        streamDebugInfo_{streamDebugInfo} {}
 
   virtual ~Decompressor() = default;
 
-  virtual uint64_t getUncompressedLength(
-      const char* /* unused */,
-      uint64_t /* unused */) const {
-    return blockSize_;
+  virtual std::pair<int64_t, bool /* Is the size exact? */>
+  getDecompressedLength(const char* /* src */, uint64_t /* srcLength */) const {
+    return {blockSize_, false};
   }
 
   virtual uint64_t decompress(
@@ -63,7 +63,7 @@ class Decompressor {
       uint64_t destLength) = 0;
 
  protected:
-  uint64_t blockSize_;
+  int64_t blockSize_;
   const std::string streamDebugInfo_;
 };
 

velox/dwio/common/compression/PagedInputStream.cpp

@@ -104,9 +104,21 @@ const char* PagedInputStream::ensureInput(size_t availableInputBytes) {
 }
 
 bool PagedInputStream::Next(const void** data, int32_t* size) {
+  VELOX_CHECK_NOT_NULL(data);
+  skipAllPending();
+  return readOrSkip(data, size);
+}
+
+// Read into `data' if it is not null; otherwise skip some of the pending.
+bool PagedInputStream::readOrSkip(const void** data, int32_t* size) {
+  if (data) {
+    VELOX_CHECK_EQ(pendingSkip_, 0);
+  }
   // if the user pushed back, return them the partial buffer
   if (outputBufferLength_) {
-    *data = outputBufferPtr_;
+    if (data) {
+      *data = outputBufferPtr_;
+    }
     *size = static_cast<int32_t>(outputBufferLength_);
     outputBufferPtr_ += outputBufferLength_;
     bytesReturned_ += outputBufferLength_;
@@ -139,7 +151,9 @@ bool PagedInputStream::Next(const void** data, int32_t* size) {
   // if no decompression or decryption is needed, simply adjust the output
   // pointer. Otherwise, make sure we have continuous block
   if (original) {
-    *data = inputBufferPtr_;
+    if (data) {
+      *data = inputBufferPtr_;
+    }
     *size = static_cast<int32_t>(availSize);
     outputBufferPtr_ = inputBufferPtr_ + availSize;
     inputBufferPtr_ += availSize;
@@ -154,24 +168,35 @@ bool PagedInputStream::Next(const void** data, int32_t* size) {
         decrypter_->decrypt(folly::StringPiece{input, remainingLength_});
     input = reinterpret_cast<const char*>(decryptionBuffer_->data());
     remainingLength_ = decryptionBuffer_->length();
-    *data = input;
+    if (data) {
+      *data = input;
+    }
     *size = remainingLength_;
     outputBufferPtr_ = input + remainingLength_;
   }
 
   // perform decompression
   if (state_ == State::START) {
     DWIO_ENSURE_NOT_NULL(decompressor_.get(), "invalid stream state");
-    prepareOutputBuffer(
-        decompressor_->getUncompressedLength(input, remainingLength_));
-    outputBufferLength_ = decompressor_->decompress(
-        input,
-        remainingLength_,
-        outputBuffer_->data(),
-        outputBuffer_->capacity());
-    *data = outputBuffer_->data();
-    *size = static_cast<int32_t>(outputBufferLength_);
-    outputBufferPtr_ = outputBuffer_->data() + outputBufferLength_;
+    DWIO_ENSURE_NOT_NULL(input);
+    auto [decompressedLength, exact] =
+        decompressor_->getDecompressedLength(input, remainingLength_);
+    if (!data && exact && decompressedLength <= pendingSkip_) {
+      *size = decompressedLength;
+      outputBufferPtr_ = nullptr;
+    } else {
+      prepareOutputBuffer(decompressedLength);
+      outputBufferLength_ = decompressor_->decompress(
+          input,
+          remainingLength_,
+          outputBuffer_->data(),
+          outputBuffer_->capacity());
+      if (data) {
+        *data = outputBuffer_->data();
+      }
+      *size = static_cast<int32_t>(outputBufferLength_);
+      outputBufferPtr_ = outputBuffer_->data() + outputBufferLength_;
+    }
     // release decryption buffer
     decryptionBuffer_ = nullptr;
   }
@@ -188,6 +213,14 @@ bool PagedInputStream::Next(const void** data, int32_t* size) {
 }
 
 void PagedInputStream::BackUp(int32_t count) {
+  if (pendingSkip_ > 0) {
+    auto len = std::min<int64_t>(count, pendingSkip_);
+    pendingSkip_ -= len;
+    count -= len;
+    if (count == 0) {
+      return;
+    }
+  }
   DWIO_ENSURE(
       outputBufferPtr_ != nullptr,
       "Backup without previous Next in ",
@@ -207,25 +240,29 @@ void PagedInputStream::BackUp(int32_t count) {
   bytesReturned_ -= count;
 }
 
-bool PagedInputStream::Skip(int32_t count) {
-  // this is a stupid implementation for now.
-  // should skip entire blocks without decompressing
-  while (count > 0) {
-    const void* ptr;
+bool PagedInputStream::skipAllPending() {
+  while (pendingSkip_ > 0) {
     int32_t len;
-    if (!Next(&ptr, &len)) {
+    if (!readOrSkip(nullptr, &len)) {
       return false;
     }
-    if (len > count) {
-      BackUp(len - count);
-      count = 0;
+    if (len > pendingSkip_) {
+      auto toBackUp = len - pendingSkip_;
+      pendingSkip_ = 0;
+      BackUp(toBackUp);
     } else {
-      count -= len;
+      pendingSkip_ -= len;
     }
   }
   return true;
 }
 
+bool PagedInputStream::Skip(int32_t count) {
+  pendingSkip_ += count;
+  // We never use the return value of this function so this is OK.
+  return true;
+}
+
 void PagedInputStream::clearDecompressionState() {
   state_ = State::HEADER;
   outputBufferLength_ = 0;
@@ -243,7 +280,8 @@ void PagedInputStream::seekToPosition(
   // to the beginning of the last view or last header, whichever is
   // later. If we are returning views into the decompression buffer,
   // we can backup to the beginning of the decompressed buffer
-  auto alreadyRead = bytesReturned_ - bytesReturnedAtLastHeaderOffset_;
+  auto alreadyRead =
+      bytesReturned_ - bytesReturnedAtLastHeaderOffset_ + pendingSkip_;
 
   // outsideOriginalWindow is true if we are returning views into
   // the input stream's buffer and we are seeking below the start of the last
@@ -260,12 +298,12 @@ void PagedInputStream::seekToPosition(
     auto provider = dwio::common::PositionProvider(positions);
     input_->seekToPosition(provider);
     clearDecompressionState();
-    Skip(uncompressedOffset);
+    pendingSkip_ = uncompressedOffset;
   } else {
     if (uncompressedOffset < alreadyRead) {
       BackUp(alreadyRead - uncompressedOffset);
     } else {
-      Skip(uncompressedOffset - alreadyRead);
+      pendingSkip_ += uncompressedOffset - alreadyRead;
     }
   }
 }