Add FSST support

Author: ArnavBalyan

cpp/src/parquet/CMakeLists.txt

@@ -171,6 +171,8 @@ set(PARQUET_SRCS
     encryption/internal_file_encryptor.cc
     exception.cc
     file_reader.cc
+    thirdparty/fsst/libfsst.cpp
+    thirdparty/fsst/fsst_avx512.cpp
     file_writer.cc
     geospatial/statistics.cc
     geospatial/util_internal.cc

cpp/src/parquet/column_reader.cc

@@ -862,7 +862,8 @@ class ColumnReaderImplBase {
         case Encoding::RLE:
         case Encoding::DELTA_BINARY_PACKED:
         case Encoding::DELTA_BYTE_ARRAY:
-        case Encoding::DELTA_LENGTH_BYTE_ARRAY: {
+        case Encoding::DELTA_LENGTH_BYTE_ARRAY:
+        case Encoding::FSST: {
           auto decoder = MakeTypedDecoder<DType>(encoding, descr_, pool_);
           current_decoder_ = decoder.get();
           decoders_[static_cast<int>(encoding)] = std::move(decoder);

cpp/src/parquet/decoder.cc

@@ -21,13 +21,15 @@
 #include <cstdint>
 #include <cstdlib>
 #include <cstring>
+#include <iostream>
 #include <limits>
 #include <memory>
 #include <string>
 #include <string_view>
 #include <type_traits>
 #include <utility>
 #include <vector>
+#include <deque>
 
 #include "arrow/array.h"
 #include "arrow/array/builder_binary.h"
@@ -51,6 +53,7 @@
 #include "parquet/exception.h"
 #include "parquet/platform.h"
 #include "parquet/schema.h"
+#include "parquet/thirdparty/fsst/fsst.h"
 #include "parquet/types.h"
 
 #ifdef _MSC_VER
@@ -2307,6 +2310,258 @@ class ByteStreamSplitDecoder<FLBAType> : public ByteStreamSplitDecoderBase<FLBAT
   }
 };
 
+// ----------------------------------------------------------------------
+
+class FsstDecoder : public DecoderImpl, virtual public TypedDecoder<ByteArrayType> {
+ public:
+  using Base = DecoderImpl;
+  using Base::num_values_;
+
+  explicit FsstDecoder(const ColumnDescriptor* descr,
+                       ::arrow::MemoryPool* pool = ::arrow::default_memory_pool())
+      : DecoderImpl(descr, Encoding::FSST), pool_(pool) {}
+
+  void SetData(int num_values, const uint8_t* data, int len) override {
+    const auto header_size = static_cast<int>(sizeof(fsst_decoder_t));
+    if (len < header_size) {
+      throw ParquetException("FSST page too small to contain decoder header");
+    }
+    num_values_ = num_values;
+    memcpy(&decoder_, data, sizeof(fsst_decoder_t));
+    next_ = data + header_size;
+    remaining_bytes_ = len - header_size;
+    decode_buffer_size_ = 0;
+  }
+
+  int Decode(ByteArray* buffer, int max_values) override {
+    max_values = std::min(max_values, num_values_);
+    if (max_values == 0) {
+      return 0;
+    }
+
+    const int64_t estimated_output_size =
+        decode_buffer_size_ +
+        static_cast<int64_t>(remaining_bytes_) * kMaxDecompressionExpansion;
+    EnsureDecodeBuffer(estimated_output_size);
+
+    int decoded = 0;
+
+    while (decoded < max_values) {
+      if (ARROW_PREDICT_FALSE(remaining_bytes_ < static_cast<int>(kLengthPrefixSize))) {
+        throw ParquetException("FSST data truncated before length prefix");
+      }
+
+      const uint32_t compressed_len = SafeLoadAs<uint32_t>(next_);
+      next_ += kLengthPrefixSize;
+      remaining_bytes_ -= static_cast<int>(kLengthPrefixSize);
+
+      if (ARROW_PREDICT_FALSE(compressed_len > static_cast<uint32_t>(remaining_bytes_))) {
+        throw ParquetException("FSST compressed length exceeds available data");
+      }
+
+      const uint8_t* compressed_ptr = next_;
+      next_ += compressed_len;
+      remaining_bytes_ -= static_cast<int>(compressed_len);
+
+      uint8_t* value_ptr = nullptr;
+      const size_t decompressed_len =
+          DecompressValue(compressed_ptr, compressed_len, &value_ptr);
+
+      buffer[decoded].ptr = value_ptr;
+      buffer[decoded].len = static_cast<uint32_t>(decompressed_len);
+      decode_buffer_size_ += static_cast<int64_t>(decompressed_len);
+      ++decoded;
+    }
+
+    num_values_ -= decoded;
+    return decoded;
+  }
+
+  int DecodeSpaced(ByteArray* buffer, int num_values, int null_count,
+                   const uint8_t* valid_bits,
+                   int64_t valid_bits_offset) override {
+    if (null_count == 0) {
+      return Decode(buffer, num_values);
+    }
+
+    const int values_to_decode = num_values - null_count;
+    temp_values_.resize(values_to_decode);
+    const int decoded = Decode(temp_values_.data(), values_to_decode);
+    if (ARROW_PREDICT_FALSE(decoded != values_to_decode)) {
+      throw ParquetException("Expected to decode ", values_to_decode,
+                             " values but decoded ", decoded, " values.");
+    }
+
+    int value_index = 0;
+    for (int i = 0; i < num_values; ++i) {
+      if (bit_util::GetBit(valid_bits, valid_bits_offset + i)) {
+        buffer[i] = temp_values_[value_index++];
+      } else {
+        buffer[i].ptr = nullptr;
+        buffer[i].len = 0;
+      }
+    }
+
+    return value_index;
+  }
+
+  int DecodeArrow(int num_values, int null_count, const uint8_t* valid_bits,
+                  int64_t valid_bits_offset,
+                  typename EncodingTraits<ByteArrayType>::Accumulator* builder) override {
+    int values_decoded = 0;
+    PARQUET_THROW_NOT_OK(
+        DecodeArrowDense(num_values, null_count, valid_bits, valid_bits_offset, builder,
+                         &values_decoded));
+    return values_decoded;
+  }
+
+  int DecodeArrow(int num_values, int null_count, const uint8_t* valid_bits,
+                  int64_t valid_bits_offset,
+                  typename EncodingTraits<ByteArrayType>::DictAccumulator* builder) override {
+    int values_decoded = 0;
+    PARQUET_THROW_NOT_OK(
+        DecodeArrowDict(num_values, null_count, valid_bits, valid_bits_offset, builder,
+                        &values_decoded));
+    return values_decoded;
+  }
+
+ private:
+  Status DecodeArrowDense(int num_values, int null_count, const uint8_t* valid_bits,
+                          int64_t valid_bits_offset,
+                          typename EncodingTraits<ByteArrayType>::Accumulator* out,
+                          int* out_values_decoded) {
+    const int values_to_decode = num_values - null_count;
+    temp_values_.resize(values_to_decode);
+
+    const int decoded = Decode(temp_values_.data(), values_to_decode);
+    if (ARROW_PREDICT_FALSE(decoded != values_to_decode)) {
+      throw ParquetException("Expected to decode ", values_to_decode,
+                             " values but decoded ", decoded, " values.");
+    }
+
+    auto visit_binary_helper = [&](auto* helper) {
+      auto* values_ptr = temp_values_.data();
+      int value_index = 0;
+
+      RETURN_NOT_OK(
+          VisitBitRuns(valid_bits, valid_bits_offset, num_values,
+                       [&](int64_t position, int64_t run_length, bool is_valid) {
+                         if (is_valid) {
+                           for (int64_t i = 0; i < run_length; ++i) {
+                             const auto& value = values_ptr[value_index++];
+                             RETURN_NOT_OK(helper->AppendValue(
+                                 value.ptr, static_cast<int32_t>(value.len)));
+                           }
+                         } else {
+                           RETURN_NOT_OK(helper->AppendNulls(run_length));
+                         }
+                         return Status::OK();
+                       }));
+
+      *out_values_decoded = decoded;
+      return Status::OK();
+    };
+
+    return DispatchArrowBinaryHelper<ByteArrayType>(
+        out, num_values, /*estimated_data_length=*/{}, visit_binary_helper);
+  }
+
+  Status DecodeArrowDict(int num_values, int null_count, const uint8_t* valid_bits,
+                         int64_t valid_bits_offset,
+                         typename EncodingTraits<ByteArrayType>::DictAccumulator* builder,
+                         int* out_values_decoded) {
+    const int values_to_decode = num_values - null_count;
+    temp_values_.resize(values_to_decode);
+
+    const int decoded = Decode(temp_values_.data(), values_to_decode);
+    if (ARROW_PREDICT_FALSE(decoded != values_to_decode)) {
+      throw ParquetException("Expected to decode ", values_to_decode,
+                             " values but decoded ", decoded, " values.");
+    }
+
+    RETURN_NOT_OK(builder->Reserve(num_values));
+
+    int value_index = 0;
+    RETURN_NOT_OK(VisitBitRuns(
+        valid_bits, valid_bits_offset, num_values,
+        [&](int64_t position, int64_t run_length, bool is_valid) {
+          if (is_valid) {
+            for (int64_t i = 0; i < run_length; ++i) {
+              const auto& value = temp_values_[value_index++];
+              RETURN_NOT_OK(builder->Append(value.ptr, static_cast<int32_t>(value.len)));
+            }
+          } else {
+            RETURN_NOT_OK(builder->AppendNulls(run_length));
+          }
+          return Status::OK();
+        }));
+
+    *out_values_decoded = decoded;
+    return Status::OK();
+  }
+
+  uint8_t* EnsureDecodeBuffer(int64_t capacity) {
+    const int64_t min_capacity =
+        std::max<int64_t>(capacity, kInitialDecodeBufferSize);
+    const int64_t target = ::arrow::bit_util::NextPower2(min_capacity);
+
+    if (!decode_buffer_) {
+      PARQUET_ASSIGN_OR_THROW(
+          decode_buffer_, ::arrow::AllocateResizableBuffer(target, pool_));
+    } else if (decode_buffer_->size() < target) {
+      PARQUET_THROW_NOT_OK(decode_buffer_->Resize(target, false));
+    }
+    return decode_buffer_->mutable_data();
+  }
+
+  size_t DecompressValue(const uint8_t* compressed_ptr, uint32_t compressed_len,
+                         uint8_t** value_ptr) {
+    EnsureDecodeBuffer(decode_buffer_size_ +
+                       OutputUpperBound(compressed_len));
+
+    while (true) {
+      uint8_t* destination = decode_buffer_->mutable_data() + decode_buffer_size_;
+      const size_t available =
+          static_cast<size_t>(decode_buffer_->size() - decode_buffer_size_);
+
+      const size_t decompressed =
+          fsst_decompress(&decoder_, compressed_len, compressed_ptr, available,
+                          destination);
+
+      if (decompressed > 0 || compressed_len == 0) {
+        *value_ptr = destination;
+        return decompressed;
+      }
+
+      int64_t new_capacity = std::max<int64_t>(
+          decode_buffer_->size() * 2,
+          decode_buffer_size_ + OutputUpperBound(compressed_len));
+      if (new_capacity <= decode_buffer_->size()) {
+        throw ParquetException("FSST decompression failed");
+      }
+      EnsureDecodeBuffer(new_capacity);
+    }
+  }
+
+  static int64_t OutputUpperBound(uint32_t compressed_len) {
+    const int64_t expanded =
+        static_cast<int64_t>(compressed_len) * kMaxDecompressionExpansion;
+    return std::max<int64_t>(expanded, kInitialDecodeBufferSize);
+  }
+
+  static constexpr size_t kLengthPrefixSize = sizeof(uint32_t);
+  static constexpr int64_t kInitialDecodeBufferSize = 1024;
+  static constexpr int64_t kMaxDecompressionExpansion = 8;
+
+  fsst_decoder_t decoder_{};
+  const uint8_t* next_ = nullptr;
+  int remaining_bytes_ = 0;
+  ::arrow::MemoryPool* pool_;
+  std::shared_ptr<::arrow::ResizableBuffer> decode_buffer_;
+  int64_t decode_buffer_size_ = 0;
+  std::vector<ByteArray> temp_values_;
+};
+
 }  // namespace
 
 // ----------------------------------------------------------------------
@@ -2373,6 +2628,13 @@ std::unique_ptr<Decoder> MakeDecoder(Type::type type_num, Encoding::type encodin
         throw ParquetException(
             "DELTA_BYTE_ARRAY only supports BYTE_ARRAY and FIXED_LEN_BYTE_ARRAY");
     }
+  } else if (encoding == Encoding::FSST) {
+    switch (type_num) {
+      case Type::BYTE_ARRAY:
+        return std::make_unique<FsstDecoder>(descr, pool);
+      default:
+        throw ParquetException("FSST encoding only supports BYTE_ARRAY");
+    }
   } else if (encoding == Encoding::DELTA_LENGTH_BYTE_ARRAY) {
     if (type_num == Type::BYTE_ARRAY) {
       return std::make_unique<DeltaLengthByteArrayDecoder>(descr, pool);