Cache ZSTD compression/decompression context #48187

Author: Andreas Ringlstetter

cpp/src/arrow/util/compression_zstd.cc

@@ -40,12 +40,27 @@ Status ZSTDError(size_t ret, const char* prefix_msg) {
   return Status::IOError(prefix_msg, ZSTD_getErrorName(ret));
 }
 
+// ----------------------------------------------------------------------
+// ZSTD context deleter implementation
+struct ZSTDContextDeleter {
+  void operator()(ZSTD_DCtx* r) {
+    if (r) {
+      ZSTD_freeDCtx(r);
+    }
+  }
+  void operator()(ZSTD_CCtx* r) {
+    if (r) {
+      ZSTD_freeCCtx(r);
+    }
+  }
+};
+
 // ----------------------------------------------------------------------
 // ZSTD decompressor implementation
 
 class ZSTDDecompressor : public Decompressor {
  public:
-  ZSTDDecompressor() : stream_(ZSTD_createDStream()) {}
+  ZSTDDecompressor() : stream_(ZSTD_createDStream()), finished_(false) {}
 
   ~ZSTDDecompressor() override { ZSTD_freeDStream(stream_); }
 
@@ -187,9 +202,13 @@ class ZSTDCodec : public Codec {
       DCHECK_EQ(output_buffer_len, 0);
       output_buffer = &empty_buffer;
     }
-
-    size_t ret = ZSTD_decompress(output_buffer, static_cast<size_t>(output_buffer_len),
-                                 input, static_cast<size_t>(input_len));
+    if (!decompression_context_) {
+      decompression_context_ =
+          decltype(decompression_context_)(ZSTD_createDCtx(), ZSTDContextDeleter{});
+    }
+    size_t ret = ZSTD_decompressDCtx(decompression_context_.get(), output_buffer,
+                                     static_cast<size_t>(output_buffer_len), input,
+                                     static_cast<size_t>(input_len));
     if (ZSTD_isError(ret)) {
       return ZSTDError(ret, "ZSTD decompression failed: ");
     }
@@ -207,8 +226,13 @@ class ZSTDCodec : public Codec {
 
   Result<int64_t> Compress(int64_t input_len, const uint8_t* input,
                            int64_t output_buffer_len, uint8_t* output_buffer) override {
-    size_t ret = ZSTD_compress(output_buffer, static_cast<size_t>(output_buffer_len),
-                               input, static_cast<size_t>(input_len), compression_level_);
+    if (!compression_context_) {
+      compression_context_ =
+          decltype(compression_context_)(ZSTD_createCCtx(), ZSTDContextDeleter{});
+    }
+    size_t ret = ZSTD_compressCCtx(compression_context_.get(), output_buffer,
+                                   static_cast<size_t>(output_buffer_len), input,
+                                   static_cast<size_t>(input_len), compression_level_);
     if (ZSTD_isError(ret)) {
       return ZSTDError(ret, "ZSTD compression failed: ");
     }
@@ -236,6 +260,8 @@ class ZSTDCodec : public Codec {
 
  private:
   const int compression_level_;
+  std::unique_ptr<ZSTD_DCtx, ZSTDContextDeleter> decompression_context_;
+  std::unique_ptr<ZSTD_CCtx, ZSTDContextDeleter> compression_context_;
 };
 
 }  // namespace