luffy-zh commented on code in PR #1916:
URL: https://github.com/apache/orc/pull/1916#discussion_r1590925121


##########
c++/src/Compression.cc:
##########
@@ -187,45 +161,81 @@ namespace orc {
 
     virtual bool Next(void** data, int* size) override;
     virtual std::string getName() const override = 0;
+    virtual void BackUp(int count) override;
+    virtual void suppress() override;
+    virtual uint64_t flush() override;
 
    protected:
     // return total compressed size
     virtual uint64_t doStreamingCompression() = 0;
+
+    // Buffer to hold uncompressed data until user calls Next()
+    BlockBuffer rawInputBuffer;
   };
 
+  void CompressionStream::BackUp(int count) {
+    uint64_t backup = static_cast<uint64_t>(count);
+    uint64_t currSize = rawInputBuffer.size();
+    if (backup > currSize) {
+      throw std::logic_error("Can't backup that much!");
+    }
+    rawInputBuffer.resize(currSize - backup);
+  }
+
+  uint64_t CompressionStream::flush() {
+    void* data;
+    int size;
+    if (!Next(&data, &size)) {
+      throw std::runtime_error("Failed to flush compression buffer.");
+    }
+    BufferedOutputStream::BackUp(outputSize - outputPosition);
+    rawInputBuffer.resize(0);
+    outputSize = outputPosition = 0;
+    return BufferedOutputStream::flush();
+  }
+
+  void CompressionStream::suppress() {
+    outputBuffer = nullptr;
+    outputPosition = outputSize = 0;
+    rawInputBuffer.resize(0);
+    BufferedOutputStream::suppress();
+  }
+
   CompressionStream::CompressionStream(OutputStream* outStream, int 
compressionLevel,
                                        uint64_t capacity, uint64_t blockSize, 
MemoryPool& pool,
                                        WriterMetrics* metrics)
-      : CompressionStreamBase(outStream, compressionLevel, capacity, 
blockSize, pool, metrics) {
+      : CompressionStreamBase(outStream, compressionLevel, capacity, 
blockSize, pool, metrics),
+        rawInputBuffer(pool, blockSize) {
     // PASS
   }
 
   bool CompressionStream::Next(void** data, int* size) {
-    if (bufferSize != 0) {
+    if (rawInputBuffer.size() != 0) {

Review Comment:
   Great advice, I'll address it in my next work project.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to