3 files changed, 56 insertions, 8 deletions

diff --git a/table/format.cc b/table/format.cc
index e183977..7647372 100644
--- a/table/format.cc
+++ b/table/format.cc
@@ -5,6 +5,7 @@
 #include "table/format.h"
 
 #include "leveldb/env.h"
+#include "leveldb/options.h"
 #include "port/port.h"
 #include "table/block.h"
 #include "util/coding.h"
@@ -116,13 +117,31 @@ Status ReadBlock(RandomAccessFile* file, const ReadOptions& options,
       size_t ulength = 0;
       if (!port::Snappy_GetUncompressedLength(data, n, &ulength)) {
         delete[] buf;
-        return Status::Corruption("corrupted compressed block contents");
+        return Status::Corruption("corrupted snappy compressed block length");
       }
       char* ubuf = new char[ulength];
       if (!port::Snappy_Uncompress(data, n, ubuf)) {
         delete[] buf;
         delete[] ubuf;
-        return Status::Corruption("corrupted compressed block contents");
+        return Status::Corruption("corrupted snappy compressed block contents");
+      }
+      delete[] buf;
+      result->data = Slice(ubuf, ulength);
+      result->heap_allocated = true;
+      result->cachable = true;
+      break;
+    }
+    case kZstdCompression: {
+      size_t ulength = 0;
+      if (!port::Zstd_GetUncompressedLength(data, n, &ulength)) {
+        delete[] buf;
+        return Status::Corruption("corrupted zstd compressed block length");
+      }
+      char* ubuf = new char[ulength];
+      if (!port::Zstd_Uncompress(data, n, ubuf)) {
+        delete[] buf;
+        delete[] ubuf;
+        return Status::Corruption("corrupted zstd compressed block contents");
       }
       delete[] buf;
       result->data = Slice(ubuf, ulength);
diff --git a/table/table_builder.cc b/table/table_builder.cc
index 29a619d..ba3df9e 100644
--- a/table/table_builder.cc
+++ b/table/table_builder.cc
@@ -168,6 +168,20 @@ void TableBuilder::WriteBlock(BlockBuilder* block, BlockHandle* handle) {
       }
       break;
     }
+
+    case kZstdCompression: {
+      std::string* compressed = &r->compressed_output;
+      if (port::Zstd_Compress(raw.data(), raw.size(), compressed) &&
+          compressed->size() < raw.size() - (raw.size() / 8u)) {
+        block_contents = *compressed;
+      } else {
+        // Zstd not supported, or compressed less than 12.5%, so just
+        // store uncompressed form
+        block_contents = raw;
+        type = kNoCompression;
+      }
+      break;
+    }
   }
   WriteRawBlock(block_contents, type, handle);
   r->compressed_output.clear();
diff --git a/table/table_test.cc b/table/table_test.cc
index a405586..b3baf95 100644
--- a/table/table_test.cc
+++ b/table/table_test.cc
@@ -14,6 +14,7 @@
 #include "leveldb/db.h"
 #include "leveldb/env.h"
 #include "leveldb/iterator.h"
+#include "leveldb/options.h"
 #include "leveldb/table_builder.h"
 #include "table/block.h"
 #include "table/block_builder.h"
@@ -784,15 +785,29 @@ TEST(TableTest, ApproximateOffsetOfPlain) {
   ASSERT_TRUE(Between(c.ApproximateOffsetOf("xyz"), 610000, 612000));
 }
 
-static bool SnappyCompressionSupported() {
+static bool CompressionSupported(CompressionType type) {
   std::string out;
   Slice in = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa";
-  return port::Snappy_Compress(in.data(), in.size(), &out);
+  if (type == kSnappyCompression) {
+    return port::Snappy_Compress(in.data(), in.size(), &out);
+  } else if (type == kZstdCompression) {
+    return port::Zstd_Compress(in.data(), in.size(), &out);
+  }
+  return false;
 }
 
-TEST(TableTest, ApproximateOffsetOfCompressed) {
-  if (!SnappyCompressionSupported())
-    GTEST_SKIP() << "skipping compression tests";
+class CompressionTableTest
+    : public ::testing::TestWithParam<std::tuple<CompressionType>> {};
+
+INSTANTIATE_TEST_SUITE_P(CompressionTests, CompressionTableTest,
+                         ::testing::Values(kSnappyCompression,
+                                           kZstdCompression));
+
+TEST_P(CompressionTableTest, ApproximateOffsetOfCompressed) {
+  CompressionType type = ::testing::get<0>(GetParam());
+  if (!CompressionSupported(type)) {
+    GTEST_SKIP() << "skipping compression test: " << type;
+  }
 
   Random rnd(301);
   TableConstructor c(BytewiseComparator());
@@ -805,7 +820,7 @@ TEST(TableTest, ApproximateOffsetOfCompressed) {
   KVMap kvmap;
   Options options;
   options.block_size = 1024;
-  options.compression = kSnappyCompression;
+  options.compression = type;
   c.Finish(options, &keys, &kvmap);
 
   // Expected upper and lower bounds of space used by compressible strings.