arrow/util/compression_lz4.cc

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include "arrow/util/compression.h"

#include <cstdint>
#include <cstring>
#include <memory>

#include <lz4.h>
#include <lz4frame.h>

#include "arrow/result.h"
#include "arrow/status.h"
#include "arrow/util/bit_util.h"
#include "arrow/util/endian.h"
#include "arrow/util/logging.h"
#include "arrow/util/macros.h"
#include "arrow/util/ubsan.h"

#ifndef LZ4F_HEADER_SIZE_MAX
#define LZ4F_HEADER_SIZE_MAX 19
#endif

namespace arrow {
namespace util {

namespace {

static Status LZ4Error(LZ4F_errorCode_t ret, const char* prefix_msg) {
  return Status::IOError(prefix_msg, LZ4F_getErrorName(ret));
}

static LZ4F_preferences_t DefaultPreferences() {
  LZ4F_preferences_t prefs;
  memset(&prefs, 0, sizeof(prefs));
  return prefs;
}

// ----------------------------------------------------------------------
// Lz4 frame decompressor implementation

class LZ4Decompressor : public Decompressor {
 public:
  LZ4Decompressor() {}

  ~LZ4Decompressor() override {
    if (ctx_ != nullptr) {
      ARROW_UNUSED(LZ4F_freeDecompressionContext(ctx_));
    }
  }

  Status Init() {
    LZ4F_errorCode_t ret;
    finished_ = false;

    ret = LZ4F_createDecompressionContext(&ctx_, LZ4F_VERSION);
    if (LZ4F_isError(ret)) {
      return LZ4Error(ret, "LZ4 init failed: ");
    } else {
      return Status::OK();
    }
  }

  Status Reset() override {
#if defined(LZ4_VERSION_NUMBER) && LZ4_VERSION_NUMBER >= 10800
    // LZ4F_resetDecompressionContext appeared in 1.8.0
    DCHECK_NE(ctx_, nullptr);
    LZ4F_resetDecompressionContext(ctx_);
    finished_ = false;
    return Status::OK();
#else
    if (ctx_ != nullptr) {
      ARROW_UNUSED(LZ4F_freeDecompressionContext(ctx_));
    }
    return Init();
#endif
  }

  Result<DecompressResult> Decompress(int64_t input_len, const uint8_t* input,
                                      int64_t output_len, uint8_t* output) override {
    auto src = input;
    auto dst = output;
    auto src_size = static_cast<size_t>(input_len);
    auto dst_capacity = static_cast<size_t>(output_len);
    size_t ret;

    ret =
        LZ4F_decompress(ctx_, dst, &dst_capacity, src, &src_size, nullptr /* options */);
    if (LZ4F_isError(ret)) {
      return LZ4Error(ret, "LZ4 decompress failed: ");
    }
    finished_ = (ret == 0);
    return DecompressResult{static_cast<int64_t>(src_size),
                            static_cast<int64_t>(dst_capacity),
                            (src_size == 0 && dst_capacity == 0)};
  }

  bool IsFinished() override { return finished_; }

 protected:
  LZ4F_decompressionContext_t ctx_ = nullptr;
  bool finished_;
};

// ----------------------------------------------------------------------
// Lz4 frame compressor implementation

class LZ4Compressor : public Compressor {
 public:
  LZ4Compressor() {}

  ~LZ4Compressor() override {
    if (ctx_ != nullptr) {
      ARROW_UNUSED(LZ4F_freeCompressionContext(ctx_));
    }
  }

  Status Init() {
    LZ4F_errorCode_t ret;
    prefs_ = DefaultPreferences();
    first_time_ = true;

    ret = LZ4F_createCompressionContext(&ctx_, LZ4F_VERSION);
    if (LZ4F_isError(ret)) {
      return LZ4Error(ret, "LZ4 init failed: ");
    } else {
      return Status::OK();
    }
  }

#define BEGIN_COMPRESS(dst, dst_capacity, output_too_small)     \
  if (first_time_) {                                            \
    if (dst_capacity < LZ4F_HEADER_SIZE_MAX) {                  \
      /* Output too small to write LZ4F header */               \
      return (output_too_small);                                \
    }                                                           \
    ret = LZ4F_compressBegin(ctx_, dst, dst_capacity, &prefs_); \
    if (LZ4F_isError(ret)) {                                    \
      return LZ4Error(ret, "LZ4 compress begin failed: ");      \
    }                                                           \
    first_time_ = false;                                        \
    dst += ret;                                                 \
    dst_capacity -= ret;                                        \
    bytes_written += static_cast<int64_t>(ret);                 \
  }

  Result<CompressResult> Compress(int64_t input_len, const uint8_t* input,
                                  int64_t output_len, uint8_t* output) override {
    auto src = input;
    auto dst = output;
    auto src_size = static_cast<size_t>(input_len);
    auto dst_capacity = static_cast<size_t>(output_len);
    size_t ret;
    int64_t bytes_written = 0;

    BEGIN_COMPRESS(dst, dst_capacity, (CompressResult{0, 0}));

    if (dst_capacity < LZ4F_compressBound(src_size, &prefs_)) {
      // Output too small to compress into
      return CompressResult{0, bytes_written};
    }
    ret = LZ4F_compressUpdate(ctx_, dst, dst_capacity, src, src_size,
                              nullptr /* options */);
    if (LZ4F_isError(ret)) {
      return LZ4Error(ret, "LZ4 compress update failed: ");
    }
    bytes_written += static_cast<int64_t>(ret);
    DCHECK_LE(bytes_written, output_len);
    return CompressResult{input_len, bytes_written};
  }

  Result<FlushResult> Flush(int64_t output_len, uint8_t* output) override {
    auto dst = output;
    auto dst_capacity = static_cast<size_t>(output_len);
    size_t ret;
    int64_t bytes_written = 0;

    BEGIN_COMPRESS(dst, dst_capacity, (FlushResult{0, true}));

    if (dst_capacity < LZ4F_compressBound(0, &prefs_)) {
      // Output too small to flush into
      return FlushResult{bytes_written, true};
    }

    ret = LZ4F_flush(ctx_, dst, dst_capacity, nullptr /* options */);
    if (LZ4F_isError(ret)) {
      return LZ4Error(ret, "LZ4 flush failed: ");
    }
    bytes_written += static_cast<int64_t>(ret);
    DCHECK_LE(bytes_written, output_len);
    return FlushResult{bytes_written, false};
  }

  Result<EndResult> End(int64_t output_len, uint8_t* output) override {
    auto dst = output;
    auto dst_capacity = static_cast<size_t>(output_len);
    size_t ret;
    int64_t bytes_written = 0;

    BEGIN_COMPRESS(dst, dst_capacity, (EndResult{0, true}));

    if (dst_capacity < LZ4F_compressBound(0, &prefs_)) {
      // Output too small to end frame into
      return EndResult{bytes_written, true};
    }

    ret = LZ4F_compressEnd(ctx_, dst, dst_capacity, nullptr /* options */);
    if (LZ4F_isError(ret)) {
      return LZ4Error(ret, "LZ4 end failed: ");
    }
    bytes_written += static_cast<int64_t>(ret);
    DCHECK_LE(bytes_written, output_len);
    return EndResult{bytes_written, false};
  }

#undef BEGIN_COMPRESS

 protected:
  LZ4F_compressionContext_t ctx_ = nullptr;
  LZ4F_preferences_t prefs_;
  bool first_time_;
};

// ----------------------------------------------------------------------
// Lz4 frame codec implementation

class Lz4FrameCodec : public Codec {
 public:
  Lz4FrameCodec() : prefs_(DefaultPreferences()) {}

  int64_t MaxCompressedLen(int64_t input_len,
                           const uint8_t* ARROW_ARG_UNUSED(input)) override {
    return static_cast<int64_t>(
        LZ4F_compressFrameBound(static_cast<size_t>(input_len), &prefs_));
  }

  Result<int64_t> Compress(int64_t input_len, const uint8_t* input,
                           int64_t output_buffer_len, uint8_t* output_buffer) override {
    auto output_len =
        LZ4F_compressFrame(output_buffer, static_cast<size_t>(output_buffer_len), input,
                           static_cast<size_t>(input_len), &prefs_);
    if (LZ4F_isError(output_len)) {
      return LZ4Error(output_len, "Lz4 compression failure: ");
    }
    return static_cast<int64_t>(output_len);
  }

  Result<int64_t> Decompress(int64_t input_len, const uint8_t* input,
                             int64_t output_buffer_len, uint8_t* output_buffer) override {
    ARROW_ASSIGN_OR_RAISE(auto decomp, MakeDecompressor());

    int64_t total_bytes_written = 0;
    while (!decomp->IsFinished() && input_len != 0) {
      ARROW_ASSIGN_OR_RAISE(
          auto res,
          decomp->Decompress(input_len, input, output_buffer_len, output_buffer));
      input += res.bytes_read;
      input_len -= res.bytes_read;
      output_buffer += res.bytes_written;
      output_buffer_len -= res.bytes_written;
      total_bytes_written += res.bytes_written;
      if (res.need_more_output) {
        return Status::IOError("Lz4 decompression buffer too small");
      }
    }
    if (!decomp->IsFinished()) {
      return Status::IOError("Lz4 compressed input contains less than one frame");
    }
    if (input_len != 0) {
      return Status::IOError("Lz4 compressed input contains more than one frame");
    }
    return total_bytes_written;
  }

  Result<std::shared_ptr<Compressor>> MakeCompressor() override {
    auto ptr = std::make_shared<LZ4Compressor>();
    RETURN_NOT_OK(ptr->Init());
    return ptr;
  }

  Result<std::shared_ptr<Decompressor>> MakeDecompressor() override {
    auto ptr = std::make_shared<LZ4Decompressor>();
    RETURN_NOT_OK(ptr->Init());
    return ptr;
  }

  Compression::type compression_type() const override { return Compression::LZ4_FRAME; }

 protected:
  const LZ4F_preferences_t prefs_;
};

// ----------------------------------------------------------------------
// Lz4 "raw" codec implementation

class Lz4Codec : public Codec {
 public:
  Result<int64_t> Decompress(int64_t input_len, const uint8_t* input,
                             int64_t output_buffer_len, uint8_t* output_buffer) override {
    int64_t decompressed_size = LZ4_decompress_safe(
        reinterpret_cast<const char*>(input), reinterpret_cast<char*>(output_buffer),
        static_cast<int>(input_len), static_cast<int>(output_buffer_len));
    if (decompressed_size < 0) {
      return Status::IOError("Corrupt Lz4 compressed data.");
    }
    return decompressed_size;
  }

  int64_t MaxCompressedLen(int64_t input_len,
                           const uint8_t* ARROW_ARG_UNUSED(input)) override {
    return LZ4_compressBound(static_cast<int>(input_len));
  }

  Result<int64_t> Compress(int64_t input_len, const uint8_t* input,
                           int64_t output_buffer_len, uint8_t* output_buffer) override {
    int64_t output_len = LZ4_compress_default(
        reinterpret_cast<const char*>(input), reinterpret_cast<char*>(output_buffer),
        static_cast<int>(input_len), static_cast<int>(output_buffer_len));
    if (output_len == 0) {
      return Status::IOError("Lz4 compression failure.");
    }
    return output_len;
  }

  Result<std::shared_ptr<Compressor>> MakeCompressor() override {
    return Status::NotImplemented(
        "Streaming compression unsupported with LZ4 raw format. "
        "Try using LZ4 frame format instead.");
  }

  Result<std::shared_ptr<Decompressor>> MakeDecompressor() override {
    return Status::NotImplemented(
        "Streaming decompression unsupported with LZ4 raw format. "
        "Try using LZ4 frame format instead.");
  }

  Compression::type compression_type() const override { return Compression::LZ4; }
};

// ----------------------------------------------------------------------
// Lz4 Hadoop "raw" codec implementation

class Lz4HadoopCodec : public Lz4Codec {
 public:
  Result<int64_t> Decompress(int64_t input_len, const uint8_t* input,
                             int64_t output_buffer_len, uint8_t* output_buffer) override {
    const int64_t decompressed_size =
        TryDecompressHadoop(input_len, input, output_buffer_len, output_buffer);
    if (decompressed_size != kNotHadoop) {
      return decompressed_size;
    }
    // Fall back on raw LZ4 codec (for files produces by earlier versions of Parquet C++)
    return Lz4Codec::Decompress(input_len, input, output_buffer_len, output_buffer);
  }

  int64_t MaxCompressedLen(int64_t input_len,
                           const uint8_t* ARROW_ARG_UNUSED(input)) override {
    return kPrefixLength + Lz4Codec::MaxCompressedLen(input_len, nullptr);
  }

  Result<int64_t> Compress(int64_t input_len, const uint8_t* input,
                           int64_t output_buffer_len, uint8_t* output_buffer) override {
    if (output_buffer_len < kPrefixLength) {
      return Status::Invalid("Output buffer too small for Lz4HadoopCodec compression");
    }

    ARROW_ASSIGN_OR_RAISE(
        int64_t output_len,
        Lz4Codec::Compress(input_len, input, output_buffer_len - kPrefixLength,
                           output_buffer + kPrefixLength));

    // Prepend decompressed size in bytes and compressed size in bytes
    // to be compatible with Hadoop Lz4Codec
    const uint32_t decompressed_size =
        BitUtil::ToBigEndian(static_cast<uint32_t>(input_len));
    const uint32_t compressed_size =
        BitUtil::ToBigEndian(static_cast<uint32_t>(output_len));
    SafeStore(output_buffer, decompressed_size);
    SafeStore(output_buffer + sizeof(uint32_t), compressed_size);

    return kPrefixLength + output_len;
  }

  Result<std::shared_ptr<Compressor>> MakeCompressor() override {
    return Status::NotImplemented(
        "Streaming compression unsupported with LZ4 Hadoop raw format. "
        "Try using LZ4 frame format instead.");
  }

  Result<std::shared_ptr<Decompressor>> MakeDecompressor() override {
    return Status::NotImplemented(
        "Streaming decompression unsupported with LZ4 Hadoop raw format. "
        "Try using LZ4 frame format instead.");
  }

  Compression::type compression_type() const override { return Compression::LZ4_HADOOP; }

 protected:
  // Offset starting at which page data can be read/written
  static const int64_t kPrefixLength = sizeof(uint32_t) * 2;

  static const int64_t kNotHadoop = -1;

  int64_t TryDecompressHadoop(int64_t input_len, const uint8_t* input,
                              int64_t output_buffer_len, uint8_t* output_buffer) {
    // Parquet files written with the Hadoop Lz4Codec use their own framing.
    // The input buffer can contain an arbitrary number of "frames", each
    // with the following structure:
    // - bytes 0..3: big-endian uint32_t representing the frame decompressed size
    // - bytes 4..7: big-endian uint32_t representing the frame compressed size
    // - bytes 8...: frame compressed data
    //
    // The Hadoop Lz4Codec source code can be found here:
    // https://github.com/apache/hadoop/blob/trunk/hadoop-mapreduce-project/hadoop-mapreduce-client/hadoop-mapreduce-client-nativetask/src/main/native/src/codec/Lz4Codec.cc
    int64_t total_decompressed_size = 0;

    while (input_len >= kPrefixLength) {
      const uint32_t expected_decompressed_size =
          BitUtil::FromBigEndian(SafeLoadAs<uint32_t>(input));
      const uint32_t expected_compressed_size =
          BitUtil::FromBigEndian(SafeLoadAs<uint32_t>(input + sizeof(uint32_t)));
      input += kPrefixLength;
      input_len -= kPrefixLength;

      if (input_len < expected_compressed_size) {
        // Not enough bytes for Hadoop "frame"
        return kNotHadoop;
      }
      if (output_buffer_len < expected_decompressed_size) {
        // Not enough bytes to hold advertised output => probably not Hadoop
        return kNotHadoop;
      }
      // Try decompressing and compare with expected decompressed length
      auto maybe_decompressed_size = Lz4Codec::Decompress(
          expected_compressed_size, input, output_buffer_len, output_buffer);
      if (!maybe_decompressed_size.ok() ||
          *maybe_decompressed_size != expected_decompressed_size) {
        return kNotHadoop;
      }
      input += expected_compressed_size;
      input_len -= expected_compressed_size;
      output_buffer += expected_decompressed_size;
      output_buffer_len -= expected_decompressed_size;
      total_decompressed_size += expected_decompressed_size;
    }

    if (input_len == 0) {
      return total_decompressed_size;
    } else {
      return kNotHadoop;
    }
  }
};

}  // namespace

namespace internal {

std::unique_ptr<Codec> MakeLz4FrameCodec() {
  return std::unique_ptr<Codec>(new Lz4FrameCodec());
}

std::unique_ptr<Codec> MakeLz4HadoopRawCodec() {
  return std::unique_ptr<Codec>(new Lz4HadoopCodec());
}

std::unique_ptr<Codec> MakeLz4RawCodec() {
  return std::unique_ptr<Codec>(new Lz4Codec());
}

}  // namespace internal

}  // namespace util
}  // namespace arrow