Back out "Optimize PagedInputStream::Skip" (#6870)

Summary:
Pull Request resolved: #6870

Original commit changeset: 07241aaf71e8

Original Phabricator Diff: D49501856

Reviewed By: Yuhta

Differential Revision: D49872683

fbshipit-source-id: d33553ce7ff603caed5befcb8d12c67a0154fba3

Author: Prasoon Telang

velox/dwio/common/compression/Compression.cpp

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

velox/dwio/common/compression/PagedInputStream.cpp

@@ -104,21 +104,9 @@ 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_) {
-    if (data) {
-      *data = outputBufferPtr_;
-    }
+    *data = outputBufferPtr_;
     *size = static_cast<int32_t>(outputBufferLength_);
     outputBufferPtr_ += outputBufferLength_;
     bytesReturned_ += outputBufferLength_;
@@ -151,9 +139,7 @@ bool PagedInputStream::readOrSkip(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) {
-    if (data) {
-      *data = inputBufferPtr_;
-    }
+    *data = inputBufferPtr_;
     *size = static_cast<int32_t>(availSize);
     outputBufferPtr_ = inputBufferPtr_ + availSize;
     inputBufferPtr_ += availSize;
@@ -168,35 +154,24 @@ bool PagedInputStream::readOrSkip(const void** data, int32_t* size) {
         decrypter_->decrypt(folly::StringPiece{input, remainingLength_});
     input = reinterpret_cast<const char*>(decryptionBuffer_->data());
     remainingLength_ = decryptionBuffer_->length();
-    if (data) {
-      *data = input;
-    }
+    *data = input;
     *size = remainingLength_;
     outputBufferPtr_ = input + remainingLength_;
   }
 
   // perform decompression
   if (state_ == State::START) {
     DWIO_ENSURE_NOT_NULL(decompressor_.get(), "invalid stream state");
-    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_;
-    }
+    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_;
     // release decryption buffer
     decryptionBuffer_ = nullptr;
   }
@@ -213,14 +188,6 @@ bool PagedInputStream::readOrSkip(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 ",
@@ -240,29 +207,25 @@ void PagedInputStream::BackUp(int32_t count) {
   bytesReturned_ -= count;
 }
 
-bool PagedInputStream::skipAllPending() {
-  while (pendingSkip_ > 0) {
+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;
     int32_t len;
-    if (!readOrSkip(nullptr, &len)) {
+    if (!Next(&ptr, &len)) {
       return false;
     }
-    if (len > pendingSkip_) {
-      auto toBackUp = len - pendingSkip_;
-      pendingSkip_ = 0;
-      BackUp(toBackUp);
+    if (len > count) {
+      BackUp(len - count);
+      count = 0;
     } else {
-      pendingSkip_ -= len;
+      count -= 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;
@@ -280,8 +243,7 @@ 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_ + pendingSkip_;
+  auto alreadyRead = bytesReturned_ - bytesReturnedAtLastHeaderOffset_;
 
   // outsideOriginalWindow is true if we are returning views into
   // the input stream's buffer and we are seeking below the start of the last
@@ -298,12 +260,12 @@ void PagedInputStream::seekToPosition(
     auto provider = dwio::common::PositionProvider(positions);
     input_->seekToPosition(provider);
     clearDecompressionState();
-    pendingSkip_ = uncompressedOffset;
+    Skip(uncompressedOffset);
   } else {
     if (uncompressedOffset < alreadyRead) {
       BackUp(alreadyRead - uncompressedOffset);
     } else {
-      pendingSkip_ += uncompressedOffset - alreadyRead;
+      Skip(uncompressedOffset - alreadyRead);
     }
   }
 }