// 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 <gtest/gtest.h>
#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <utility>
#include <vector>

#include "arrow/array.h"
#include "arrow/array/builder_dict.h"
#include "arrow/testing/gtest_util.h"
#include "arrow/testing/random.h"
#include "arrow/testing/util.h"
#include "arrow/type.h"
#include "arrow/util/bit_util.h"
#include "arrow/util/bitmap_writer.h"
#include "arrow/util/checked_cast.h"
#include "arrow/util/endian.h"

#include "parquet/encoding.h"
#include "parquet/platform.h"
#include "parquet/schema.h"
#include "parquet/test_util.h"
#include "parquet/types.h"

using arrow::default_memory_pool;
using arrow::MemoryPool;
using arrow::internal::checked_cast;

namespace BitUtil = arrow::BitUtil;

namespace parquet {

namespace test {

TEST(VectorBooleanTest, TestEncodeDecode) {
  // PARQUET-454
  int nvalues = 10000;
  int nbytes = static_cast<int>(BitUtil::BytesForBits(nvalues));

  std::vector<bool> draws;
  ::arrow::random_is_valid(nvalues, 0.5 /* null prob */, &draws, 0 /* seed */);

  std::unique_ptr<BooleanEncoder> encoder =
      MakeTypedEncoder<BooleanType>(Encoding::PLAIN);
  encoder->Put(draws, nvalues);

  std::unique_ptr<BooleanDecoder> decoder =
      MakeTypedDecoder<BooleanType>(Encoding::PLAIN);

  std::shared_ptr<Buffer> encode_buffer = encoder->FlushValues();
  ASSERT_EQ(nbytes, encode_buffer->size());

  std::vector<uint8_t> decode_buffer(nbytes);
  const uint8_t* decode_data = &decode_buffer[0];

  decoder->SetData(nvalues, encode_buffer->data(),
                   static_cast<int>(encode_buffer->size()));
  int values_decoded = decoder->Decode(&decode_buffer[0], nvalues);
  ASSERT_EQ(nvalues, values_decoded);

  for (int i = 0; i < nvalues; ++i) {
    ASSERT_EQ(draws[i], ::arrow::BitUtil::GetBit(decode_data, i)) << i;
  }
}

// ----------------------------------------------------------------------
// test data generation

template <typename T>
void GenerateData(int num_values, T* out, std::vector<uint8_t>* heap) {
  // seed the prng so failure is deterministic
  random_numbers(num_values, 0, std::numeric_limits<T>::min(),
                 std::numeric_limits<T>::max(), out);
}

template <>
void GenerateData<bool>(int num_values, bool* out, std::vector<uint8_t>* heap) {
  // seed the prng so failure is deterministic
  random_bools(num_values, 0.5, 0, out);
}

template <>
void GenerateData<Int96>(int num_values, Int96* out, std::vector<uint8_t>* heap) {
  // seed the prng so failure is deterministic
  random_Int96_numbers(num_values, 0, std::numeric_limits<int32_t>::min(),
                       std::numeric_limits<int32_t>::max(), out);
}

template <>
void GenerateData<ByteArray>(int num_values, ByteArray* out, std::vector<uint8_t>* heap) {
  // seed the prng so failure is deterministic
  int max_byte_array_len = 12;
  heap->resize(num_values * max_byte_array_len);
  random_byte_array(num_values, 0, heap->data(), out, 2, max_byte_array_len);
}

static int flba_length = 8;

template <>
void GenerateData<FLBA>(int num_values, FLBA* out, std::vector<uint8_t>* heap) {
  // seed the prng so failure is deterministic
  heap->resize(num_values * flba_length);
  random_fixed_byte_array(num_values, 0, heap->data(), flba_length, out);
}

template <typename T>
void VerifyResults(T* result, T* expected, int num_values) {
  for (int i = 0; i < num_values; ++i) {
    ASSERT_EQ(expected[i], result[i]) << i;
  }
}

template <typename T>
void VerifyResultsSpaced(T* result, T* expected, int num_values,
                         const uint8_t* valid_bits, int64_t valid_bits_offset) {
  for (auto i = 0; i < num_values; ++i) {
    if (BitUtil::GetBit(valid_bits, valid_bits_offset + i)) {
      ASSERT_EQ(expected[i], result[i]) << i;
    }
  }
}

template <>
void VerifyResults<FLBA>(FLBA* result, FLBA* expected, int num_values) {
  for (int i = 0; i < num_values; ++i) {
    ASSERT_EQ(0, memcmp(expected[i].ptr, result[i].ptr, flba_length)) << i;
  }
}

template <>
void VerifyResultsSpaced<FLBA>(FLBA* result, FLBA* expected, int num_values,
                               const uint8_t* valid_bits, int64_t valid_bits_offset) {
  for (auto i = 0; i < num_values; ++i) {
    if (BitUtil::GetBit(valid_bits, valid_bits_offset + i)) {
      ASSERT_EQ(0, memcmp(expected[i].ptr, result[i].ptr, flba_length)) << i;
    }
  }
}

// ----------------------------------------------------------------------
// Create some column descriptors

template <typename DType>
std::shared_ptr<ColumnDescriptor> ExampleDescr() {
  auto node = schema::PrimitiveNode::Make("name", Repetition::OPTIONAL, DType::type_num);
  return std::make_shared<ColumnDescriptor>(node, 0, 0);
}

template <>
std::shared_ptr<ColumnDescriptor> ExampleDescr<FLBAType>() {
  auto node = schema::PrimitiveNode::Make("name", Repetition::OPTIONAL,
                                          Type::FIXED_LEN_BYTE_ARRAY,
                                          ConvertedType::DECIMAL, flba_length, 10, 2);
  return std::make_shared<ColumnDescriptor>(node, 0, 0);
}

// ----------------------------------------------------------------------
// Plain encoding tests

template <typename Type>
class TestEncodingBase : public ::testing::Test {
 public:
  using c_type = typename Type::c_type;
  static constexpr int TYPE = Type::type_num;

  void SetUp() {
    descr_ = ExampleDescr<Type>();
    type_length_ = descr_->type_length();
    allocator_ = default_memory_pool();
  }

  void TearDown() {}

  void InitData(int nvalues, int repeats) {
    num_values_ = nvalues * repeats;
    input_bytes_.resize(num_values_ * sizeof(c_type));
    output_bytes_.resize(num_values_ * sizeof(c_type));
    draws_ = reinterpret_cast<c_type*>(input_bytes_.data());
    decode_buf_ = reinterpret_cast<c_type*>(output_bytes_.data());
    GenerateData<c_type>(nvalues, draws_, &data_buffer_);

    // add some repeated values
    for (int j = 1; j < repeats; ++j) {
      for (int i = 0; i < nvalues; ++i) {
        draws_[nvalues * j + i] = draws_[i];
      }
    }
  }

  virtual void CheckRoundtrip() = 0;

  virtual void CheckRoundtripSpaced(const uint8_t* valid_bits,
                                    int64_t valid_bits_offset) {}

  void Execute(int nvalues, int repeats) {
    InitData(nvalues, repeats);
    CheckRoundtrip();
  }

  void ExecuteSpaced(int nvalues, int repeats, int64_t valid_bits_offset,
                     double null_probability) {
    InitData(nvalues, repeats);

    int64_t size = num_values_ + valid_bits_offset;
    auto rand = ::arrow::random::RandomArrayGenerator(1923);
    const auto array = rand.UInt8(size, 0, 100, null_probability);
    const auto valid_bits = array->null_bitmap_data();
    if (valid_bits) {
      CheckRoundtripSpaced(valid_bits, valid_bits_offset);
    }
  }

 protected:
  MemoryPool* allocator_;

  int num_values_;
  int type_length_;
  c_type* draws_;
  c_type* decode_buf_;
  std::vector<uint8_t> input_bytes_;
  std::vector<uint8_t> output_bytes_;
  std::vector<uint8_t> data_buffer_;

  std::shared_ptr<Buffer> encode_buffer_;
  std::shared_ptr<ColumnDescriptor> descr_;
};

// Member variables are not visible to templated subclasses. Possibly figure
// out an alternative to this class layering at some point
#define USING_BASE_MEMBERS()                    \
  using TestEncodingBase<Type>::allocator_;     \
  using TestEncodingBase<Type>::descr_;         \
  using TestEncodingBase<Type>::num_values_;    \
  using TestEncodingBase<Type>::draws_;         \
  using TestEncodingBase<Type>::data_buffer_;   \
  using TestEncodingBase<Type>::type_length_;   \
  using TestEncodingBase<Type>::encode_buffer_; \
  using TestEncodingBase<Type>::decode_buf_;

template <typename Type>
class TestPlainEncoding : public TestEncodingBase<Type> {
 public:
  using c_type = typename Type::c_type;
  static constexpr int TYPE = Type::type_num;

  virtual void CheckRoundtrip() {
    auto encoder = MakeTypedEncoder<Type>(Encoding::PLAIN, false, descr_.get());
    auto decoder = MakeTypedDecoder<Type>(Encoding::PLAIN, descr_.get());
    encoder->Put(draws_, num_values_);
    encode_buffer_ = encoder->FlushValues();

    decoder->SetData(num_values_, encode_buffer_->data(),
                     static_cast<int>(encode_buffer_->size()));
    int values_decoded = decoder->Decode(decode_buf_, num_values_);
    ASSERT_EQ(num_values_, values_decoded);
    ASSERT_NO_FATAL_FAILURE(VerifyResults<c_type>(decode_buf_, draws_, num_values_));
  }

  void CheckRoundtripSpaced(const uint8_t* valid_bits, int64_t valid_bits_offset) {
    auto encoder = MakeTypedEncoder<Type>(Encoding::PLAIN, false, descr_.get());
    auto decoder = MakeTypedDecoder<Type>(Encoding::PLAIN, descr_.get());
    int null_count = 0;
    for (auto i = 0; i < num_values_; i++) {
      if (!BitUtil::GetBit(valid_bits, valid_bits_offset + i)) {
        null_count++;
      }
    }

    encoder->PutSpaced(draws_, num_values_, valid_bits, valid_bits_offset);
    encode_buffer_ = encoder->FlushValues();
    decoder->SetData(num_values_ - null_count, encode_buffer_->data(),
                     static_cast<int>(encode_buffer_->size()));
    auto values_decoded = decoder->DecodeSpaced(decode_buf_, num_values_, null_count,
                                                valid_bits, valid_bits_offset);
    ASSERT_EQ(num_values_, values_decoded);
    ASSERT_NO_FATAL_FAILURE(VerifyResultsSpaced<c_type>(decode_buf_, draws_, num_values_,
                                                        valid_bits, valid_bits_offset));
  }

 protected:
  USING_BASE_MEMBERS();
};

TYPED_TEST_SUITE(TestPlainEncoding, ParquetTypes);

TYPED_TEST(TestPlainEncoding, BasicRoundTrip) {
  ASSERT_NO_FATAL_FAILURE(this->Execute(10000, 1));

  // Spaced test with different sizes and offest to guarantee SIMD implementation
  constexpr int kAvx512Size = 64;         // sizeof(__m512i) for Avx512
  constexpr int kSimdSize = kAvx512Size;  // Current the max is Avx512
  constexpr int kMultiSimdSize = kSimdSize * 33;

  for (auto null_prob : {0.001, 0.1, 0.5, 0.9, 0.999}) {
    // Test with both size and offset up to 3 Simd block
    for (auto i = 1; i < kSimdSize * 3; i++) {
      ASSERT_NO_FATAL_FAILURE(this->ExecuteSpaced(i, 1, 0, null_prob));
      ASSERT_NO_FATAL_FAILURE(this->ExecuteSpaced(i, 1, i + 1, null_prob));
    }
    // Large block and offset
    ASSERT_NO_FATAL_FAILURE(this->ExecuteSpaced(kMultiSimdSize, 1, 0, null_prob));
    ASSERT_NO_FATAL_FAILURE(this->ExecuteSpaced(kMultiSimdSize + 33, 1, 0, null_prob));
    ASSERT_NO_FATAL_FAILURE(this->ExecuteSpaced(kMultiSimdSize, 1, 33, null_prob));
    ASSERT_NO_FATAL_FAILURE(this->ExecuteSpaced(kMultiSimdSize + 33, 1, 33, null_prob));
  }
}

// ----------------------------------------------------------------------
// Dictionary encoding tests

typedef ::testing::Types<Int32Type, Int64Type, Int96Type, FloatType, DoubleType,
                         ByteArrayType, FLBAType>
    DictEncodedTypes;

template <typename Type>
class TestDictionaryEncoding : public TestEncodingBase<Type> {
 public:
  using c_type = typename Type::c_type;
  static constexpr int TYPE = Type::type_num;

  void CheckRoundtrip() {
    std::vector<uint8_t> valid_bits(::arrow::BitUtil::BytesForBits(num_values_) + 1, 255);

    auto base_encoder = MakeEncoder(Type::type_num, Encoding::PLAIN, true, descr_.get());
    auto encoder =
        dynamic_cast<typename EncodingTraits<Type>::Encoder*>(base_encoder.get());
    auto dict_traits = dynamic_cast<DictEncoder<Type>*>(base_encoder.get());

    ASSERT_NO_THROW(encoder->Put(draws_, num_values_));
    dict_buffer_ =
        AllocateBuffer(default_memory_pool(), dict_traits->dict_encoded_size());
    dict_traits->WriteDict(dict_buffer_->mutable_data());
    std::shared_ptr<Buffer> indices = encoder->FlushValues();

    auto base_spaced_encoder =
        MakeEncoder(Type::type_num, Encoding::PLAIN, true, descr_.get());
    auto spaced_encoder =
        dynamic_cast<typename EncodingTraits<Type>::Encoder*>(base_spaced_encoder.get());

    // PutSpaced should lead to the same results
    // This also checks the PutSpaced implementation for valid_bits=nullptr
    ASSERT_NO_THROW(spaced_encoder->PutSpaced(draws_, num_values_, nullptr, 0));
    std::shared_ptr<Buffer> indices_from_spaced = spaced_encoder->FlushValues();
    ASSERT_TRUE(indices_from_spaced->Equals(*indices));

    auto dict_decoder = MakeTypedDecoder<Type>(Encoding::PLAIN, descr_.get());
    dict_decoder->SetData(dict_traits->num_entries(), dict_buffer_->data(),
                          static_cast<int>(dict_buffer_->size()));

    auto decoder = MakeDictDecoder<Type>(descr_.get());
    decoder->SetDict(dict_decoder.get());

    decoder->SetData(num_values_, indices->data(), static_cast<int>(indices->size()));
    int values_decoded = decoder->Decode(decode_buf_, num_values_);
    ASSERT_EQ(num_values_, values_decoded);

    // TODO(wesm): The DictionaryDecoder must stay alive because the decoded
    // values' data is owned by a buffer inside the DictionaryEncoder. We
    // should revisit when data lifetime is reviewed more generally.
    ASSERT_NO_FATAL_FAILURE(VerifyResults<c_type>(decode_buf_, draws_, num_values_));

    // Also test spaced decoding
    decoder->SetData(num_values_, indices->data(), static_cast<int>(indices->size()));
    // Also tests DecodeSpaced handling for valid_bits=nullptr
    values_decoded = decoder->DecodeSpaced(decode_buf_, num_values_, 0, nullptr, 0);
    ASSERT_EQ(num_values_, values_decoded);
    ASSERT_NO_FATAL_FAILURE(VerifyResults<c_type>(decode_buf_, draws_, num_values_));
  }

 protected:
  USING_BASE_MEMBERS();
  std::shared_ptr<ResizableBuffer> dict_buffer_;
};

TYPED_TEST_SUITE(TestDictionaryEncoding, DictEncodedTypes);

TYPED_TEST(TestDictionaryEncoding, BasicRoundTrip) {
  ASSERT_NO_FATAL_FAILURE(this->Execute(2500, 2));
}

TEST(TestDictionaryEncoding, CannotDictDecodeBoolean) {
  ASSERT_THROW(MakeDictDecoder<BooleanType>(nullptr), ParquetException);
}

// ----------------------------------------------------------------------
// Shared arrow builder decode tests

class TestArrowBuilderDecoding : public ::testing::Test {
 public:
  using DenseBuilder = ::arrow::internal::ChunkedBinaryBuilder;
  using DictBuilder = ::arrow::BinaryDictionary32Builder;

  void SetUp() override { null_probabilities_ = {0.0, 0.5, 1.0}; }
  void TearDown() override {}

  void InitTestCase(double null_probability) {
    GenerateInputData(null_probability);
    SetupEncoderDecoder();
  }

  void GenerateInputData(double null_probability) {
    constexpr int num_unique = 100;
    constexpr int repeat = 100;
    constexpr int64_t min_length = 2;
    constexpr int64_t max_length = 10;
    ::arrow::random::RandomArrayGenerator rag(0);
    expected_dense_ = rag.BinaryWithRepeats(repeat * num_unique, num_unique, min_length,
                                            max_length, null_probability);

    num_values_ = static_cast<int>(expected_dense_->length());
    null_count_ = static_cast<int>(expected_dense_->null_count());
    valid_bits_ = expected_dense_->null_bitmap_data();

    auto builder = CreateDictBuilder();
    ASSERT_OK(builder->AppendArray(*expected_dense_));
    ASSERT_OK(builder->Finish(&expected_dict_));

    // Initialize input_data_ for the encoder from the expected_array_ values
    const auto& binary_array = static_cast<const ::arrow::BinaryArray&>(*expected_dense_);
    input_data_.resize(binary_array.length());

    for (int64_t i = 0; i < binary_array.length(); ++i) {
      auto view = binary_array.GetView(i);
      input_data_[i] = {static_cast<uint32_t>(view.length()),
                        reinterpret_cast<const uint8_t*>(view.data())};
    }
  }

  std::unique_ptr<DictBuilder> CreateDictBuilder() {
    return std::unique_ptr<DictBuilder>(new DictBuilder(default_memory_pool()));
  }

  // Setup encoder/decoder pair for testing with
  virtual void SetupEncoderDecoder() = 0;

  void CheckDense(int actual_num_values, const ::arrow::Array& chunk) {
    ASSERT_EQ(actual_num_values, num_values_ - null_count_);
    ASSERT_ARRAYS_EQUAL(chunk, *expected_dense_);
  }

  template <typename Builder>
  void CheckDict(int actual_num_values, Builder& builder) {
    ASSERT_EQ(actual_num_values, num_values_ - null_count_);
    std::shared_ptr<::arrow::Array> actual;
    ASSERT_OK(builder.Finish(&actual));
    ASSERT_ARRAYS_EQUAL(*actual, *expected_dict_);
  }

  void CheckDecodeArrowUsingDenseBuilder() {
    for (auto np : null_probabilities_) {
      InitTestCase(np);

      typename EncodingTraits<ByteArrayType>::Accumulator acc;
      acc.builder.reset(new ::arrow::BinaryBuilder);
      auto actual_num_values =
          decoder_->DecodeArrow(num_values_, null_count_, valid_bits_, 0, &acc);

      std::shared_ptr<::arrow::Array> chunk;
      ASSERT_OK(acc.builder->Finish(&chunk));
      CheckDense(actual_num_values, *chunk);
    }
  }

  void CheckDecodeArrowUsingDictBuilder() {
    for (auto np : null_probabilities_) {
      InitTestCase(np);
      auto builder = CreateDictBuilder();
      auto actual_num_values =
          decoder_->DecodeArrow(num_values_, null_count_, valid_bits_, 0, builder.get());
      CheckDict(actual_num_values, *builder);
    }
  }

  void CheckDecodeArrowNonNullUsingDenseBuilder() {
    for (auto np : null_probabilities_) {
      InitTestCase(np);
      if (null_count_ > 0) {
        continue;
      }
      typename EncodingTraits<ByteArrayType>::Accumulator acc;
      acc.builder.reset(new ::arrow::BinaryBuilder);
      auto actual_num_values = decoder_->DecodeArrowNonNull(num_values_, &acc);
      std::shared_ptr<::arrow::Array> chunk;
      ASSERT_OK(acc.builder->Finish(&chunk));
      CheckDense(actual_num_values, *chunk);
    }
  }

  void CheckDecodeArrowNonNullUsingDictBuilder() {
    for (auto np : null_probabilities_) {
      InitTestCase(np);
      if (null_count_ > 0) {
        continue;
      }
      auto builder = CreateDictBuilder();
      auto actual_num_values = decoder_->DecodeArrowNonNull(num_values_, builder.get());
      CheckDict(actual_num_values, *builder);
    }
  }

 protected:
  std::vector<double> null_probabilities_;
  std::shared_ptr<::arrow::Array> expected_dict_;
  std::shared_ptr<::arrow::Array> expected_dense_;
  int num_values_;
  int null_count_;
  std::vector<ByteArray> input_data_;
  const uint8_t* valid_bits_;
  std::unique_ptr<ByteArrayEncoder> encoder_;
  ByteArrayDecoder* decoder_;
  std::unique_ptr<ByteArrayDecoder> plain_decoder_;
  std::unique_ptr<DictDecoder<ByteArrayType>> dict_decoder_;
  std::shared_ptr<Buffer> buffer_;
};

class PlainEncoding : public TestArrowBuilderDecoding {
 public:
  void SetupEncoderDecoder() override {
    encoder_ = MakeTypedEncoder<ByteArrayType>(Encoding::PLAIN);
    plain_decoder_ = MakeTypedDecoder<ByteArrayType>(Encoding::PLAIN);
    decoder_ = plain_decoder_.get();
    if (valid_bits_ != nullptr) {
      ASSERT_NO_THROW(
          encoder_->PutSpaced(input_data_.data(), num_values_, valid_bits_, 0));
    } else {
      ASSERT_NO_THROW(encoder_->Put(input_data_.data(), num_values_));
    }
    buffer_ = encoder_->FlushValues();
    decoder_->SetData(num_values_, buffer_->data(), static_cast<int>(buffer_->size()));
  }
};

TEST_F(PlainEncoding, CheckDecodeArrowUsingDenseBuilder) {
  this->CheckDecodeArrowUsingDenseBuilder();
}

TEST_F(PlainEncoding, CheckDecodeArrowUsingDictBuilder) {
  this->CheckDecodeArrowUsingDictBuilder();
}

TEST_F(PlainEncoding, CheckDecodeArrowNonNullDenseBuilder) {
  this->CheckDecodeArrowNonNullUsingDenseBuilder();
}

TEST_F(PlainEncoding, CheckDecodeArrowNonNullDictBuilder) {
  this->CheckDecodeArrowNonNullUsingDictBuilder();
}

TEST(PlainEncodingAdHoc, ArrowBinaryDirectPut) {
  // Implemented as part of ARROW-3246

  const int64_t size = 50;
  const int32_t min_length = 0;
  const int32_t max_length = 10;
  const double null_probability = 0.25;

  auto CheckSeed = [&](int seed) {
    ::arrow::random::RandomArrayGenerator rag(seed);
    auto values = rag.String(size, min_length, max_length, null_probability);

    auto encoder = MakeTypedEncoder<ByteArrayType>(Encoding::PLAIN);
    auto decoder = MakeTypedDecoder<ByteArrayType>(Encoding::PLAIN);

    ASSERT_NO_THROW(encoder->Put(*values));
    auto buf = encoder->FlushValues();

    int num_values = static_cast<int>(values->length() - values->null_count());
    decoder->SetData(num_values, buf->data(), static_cast<int>(buf->size()));

    typename EncodingTraits<ByteArrayType>::Accumulator acc;
    acc.builder.reset(new ::arrow::StringBuilder);
    ASSERT_EQ(num_values,
              decoder->DecodeArrow(static_cast<int>(values->length()),
                                   static_cast<int>(values->null_count()),
                                   values->null_bitmap_data(), values->offset(), &acc));

    std::shared_ptr<::arrow::Array> result;
    ASSERT_OK(acc.builder->Finish(&result));
    ASSERT_EQ(50, result->length());
    ::arrow::AssertArraysEqual(*values, *result);
  };

  for (auto seed : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}) {
    CheckSeed(seed);
  }
}

template <typename T>
void GetDictDecoder(DictEncoder<T>* encoder, int64_t num_values,
                    std::shared_ptr<Buffer>* out_values,
                    std::shared_ptr<Buffer>* out_dict, const ColumnDescriptor* descr,
                    std::unique_ptr<TypedDecoder<T>>* out_decoder) {
  auto decoder = MakeDictDecoder<T>(descr);
  auto buf = encoder->FlushValues();
  auto dict_buf = AllocateBuffer(default_memory_pool(), encoder->dict_encoded_size());
  encoder->WriteDict(dict_buf->mutable_data());

  auto dict_decoder = MakeTypedDecoder<T>(Encoding::PLAIN, descr);
  dict_decoder->SetData(encoder->num_entries(), dict_buf->data(),
                        static_cast<int>(dict_buf->size()));

  decoder->SetData(static_cast<int>(num_values), buf->data(),
                   static_cast<int>(buf->size()));
  decoder->SetDict(dict_decoder.get());

  *out_values = buf;
  *out_dict = dict_buf;
  ASSERT_NE(decoder, nullptr);
  auto released = dynamic_cast<TypedDecoder<T>*>(decoder.release());
  ASSERT_NE(released, nullptr);
  *out_decoder = std::unique_ptr<TypedDecoder<T>>(released);
}

template <typename ParquetType>
class EncodingAdHocTyped : public ::testing::Test {
 public:
  using ArrowType = typename EncodingTraits<ParquetType>::ArrowType;
  using EncoderType = typename EncodingTraits<ParquetType>::Encoder;
  using DecoderType = typename EncodingTraits<ParquetType>::Decoder;
  using BuilderType = typename EncodingTraits<ParquetType>::Accumulator;
  using DictBuilderType = typename EncodingTraits<ParquetType>::DictAccumulator;

  static const ColumnDescriptor* column_descr() {
    static auto column_descr = ExampleDescr<ParquetType>();
    return column_descr.get();
  }

  std::shared_ptr<::arrow::Array> GetValues(int seed);

  static std::shared_ptr<::arrow::DataType> arrow_type();

  void Plain(int seed) {
    auto values = GetValues(seed);
    auto encoder = MakeTypedEncoder<ParquetType>(
        Encoding::PLAIN, /*use_dictionary=*/false, column_descr());
    auto decoder = MakeTypedDecoder<ParquetType>(Encoding::PLAIN, column_descr());

    ASSERT_NO_THROW(encoder->Put(*values));
    auto buf = encoder->FlushValues();

    int num_values = static_cast<int>(values->length() - values->null_count());
    decoder->SetData(num_values, buf->data(), static_cast<int>(buf->size()));

    BuilderType acc(arrow_type(), ::arrow::default_memory_pool());
    ASSERT_EQ(num_values,
              decoder->DecodeArrow(static_cast<int>(values->length()),
                                   static_cast<int>(values->null_count()),
                                   values->null_bitmap_data(), values->offset(), &acc));

    std::shared_ptr<::arrow::Array> result;
    ASSERT_OK(acc.Finish(&result));
    ASSERT_EQ(50, result->length());
    ::arrow::AssertArraysEqual(*values, *result, /*verbose=*/true);
  }

  void ByteStreamSplit(int seed) {
    if (!std::is_same<ParquetType, FloatType>::value &&
        !std::is_same<ParquetType, DoubleType>::value) {
      return;
    }
    auto values = GetValues(seed);
    auto encoder = MakeTypedEncoder<ParquetType>(
        Encoding::BYTE_STREAM_SPLIT, /*use_dictionary=*/false, column_descr());
    auto decoder =
        MakeTypedDecoder<ParquetType>(Encoding::BYTE_STREAM_SPLIT, column_descr());

    ASSERT_NO_THROW(encoder->Put(*values));
    auto buf = encoder->FlushValues();

    int num_values = static_cast<int>(values->length() - values->null_count());
    decoder->SetData(num_values, buf->data(), static_cast<int>(buf->size()));

    BuilderType acc(arrow_type(), ::arrow::default_memory_pool());
    ASSERT_EQ(num_values,
              decoder->DecodeArrow(static_cast<int>(values->length()),
                                   static_cast<int>(values->null_count()),
                                   values->null_bitmap_data(), values->offset(), &acc));

    std::shared_ptr<::arrow::Array> result;
    ASSERT_OK(acc.Finish(&result));
    ASSERT_EQ(50, result->length());
    ::arrow::AssertArraysEqual(*values, *result);
  }

  void Dict(int seed) {
    if (std::is_same<ParquetType, BooleanType>::value) {
      return;
    }

    auto values = GetValues(seed);

    auto owned_encoder =
        MakeTypedEncoder<ParquetType>(Encoding::PLAIN,
                                      /*use_dictionary=*/true, column_descr());
    auto encoder = dynamic_cast<DictEncoder<ParquetType>*>(owned_encoder.get());

    ASSERT_NO_THROW(encoder->Put(*values));

    std::shared_ptr<Buffer> buf, dict_buf;
    int num_values = static_cast<int>(values->length() - values->null_count());

    std::unique_ptr<TypedDecoder<ParquetType>> decoder;
    GetDictDecoder(encoder, num_values, &buf, &dict_buf, column_descr(), &decoder);

    BuilderType acc(arrow_type(), ::arrow::default_memory_pool());
    ASSERT_EQ(num_values,
              decoder->DecodeArrow(static_cast<int>(values->length()),
                                   static_cast<int>(values->null_count()),
                                   values->null_bitmap_data(), values->offset(), &acc));

    std::shared_ptr<::arrow::Array> result;
    ASSERT_OK(acc.Finish(&result));
    ::arrow::AssertArraysEqual(*values, *result);
  }

  void DictPutIndices() {
    if (std::is_same<ParquetType, BooleanType>::value) {
      return;
    }

    auto dict_values = ::arrow::ArrayFromJSON(
        arrow_type(), std::is_same<ParquetType, FLBAType>::value
                          ? R"(["abcdefgh", "ijklmnop", "qrstuvwx"])"
                          : "[120, -37, 47]");
    auto indices = ::arrow::ArrayFromJSON(::arrow::int32(), "[0, 1, 2]");
    auto indices_nulls =
        ::arrow::ArrayFromJSON(::arrow::int32(), "[null, 0, 1, null, 2]");

    auto expected = ::arrow::ArrayFromJSON(
        arrow_type(), std::is_same<ParquetType, FLBAType>::value
                          ? R"(["abcdefgh", "ijklmnop", "qrstuvwx", null,
                                "abcdefgh", "ijklmnop", null, "qrstuvwx"])"
                          : "[120, -37, 47, null, "
                            "120, -37, null, 47]");

    auto owned_encoder =
        MakeTypedEncoder<ParquetType>(Encoding::PLAIN,
                                      /*use_dictionary=*/true, column_descr());
    auto owned_decoder = MakeDictDecoder<ParquetType>();

    auto encoder = dynamic_cast<DictEncoder<ParquetType>*>(owned_encoder.get());

    ASSERT_NO_THROW(encoder->PutDictionary(*dict_values));

    // Trying to call PutDictionary again throws
    ASSERT_THROW(encoder->PutDictionary(*dict_values), ParquetException);

    ASSERT_NO_THROW(encoder->PutIndices(*indices));
    ASSERT_NO_THROW(encoder->PutIndices(*indices_nulls));

    std::shared_ptr<Buffer> buf, dict_buf;
    int num_values = static_cast<int>(expected->length() - expected->null_count());

    std::unique_ptr<TypedDecoder<ParquetType>> decoder;
    GetDictDecoder(encoder, num_values, &buf, &dict_buf, column_descr(), &decoder);

    BuilderType acc(arrow_type(), ::arrow::default_memory_pool());
    ASSERT_EQ(num_values, decoder->DecodeArrow(static_cast<int>(expected->length()),
                                               static_cast<int>(expected->null_count()),
                                               expected->null_bitmap_data(),
                                               expected->offset(), &acc));

    std::shared_ptr<::arrow::Array> result;
    ASSERT_OK(acc.Finish(&result));
    ::arrow::AssertArraysEqual(*expected, *result);
  }

 protected:
  const int64_t size_ = 50;
  const double null_probability_ = 0.25;
};

template <typename ParquetType>
std::shared_ptr<::arrow::DataType> EncodingAdHocTyped<ParquetType>::arrow_type() {
  return ::arrow::TypeTraits<ArrowType>::type_singleton();
}

template <>
std::shared_ptr<::arrow::DataType> EncodingAdHocTyped<FLBAType>::arrow_type() {
  return ::arrow::fixed_size_binary(sizeof(uint64_t));
}

template <typename ParquetType>
std::shared_ptr<::arrow::Array> EncodingAdHocTyped<ParquetType>::GetValues(int seed) {
  ::arrow::random::RandomArrayGenerator rag(seed);
  return rag.Numeric<ArrowType>(size_, 0, 10, null_probability_);
}

template <>
std::shared_ptr<::arrow::Array> EncodingAdHocTyped<BooleanType>::GetValues(int seed) {
  ::arrow::random::RandomArrayGenerator rag(seed);
  return rag.Boolean(size_, 0.1, null_probability_);
}

template <>
std::shared_ptr<::arrow::Array> EncodingAdHocTyped<FLBAType>::GetValues(int seed) {
  ::arrow::random::RandomArrayGenerator rag(seed);
  std::shared_ptr<::arrow::Array> values;
  ARROW_EXPECT_OK(
      rag.UInt64(size_, 0, std::numeric_limits<uint64_t>::max(), null_probability_)
          ->View(arrow_type())
          .Value(&values));
  return values;
}

using EncodingAdHocTypedCases =
    ::testing::Types<BooleanType, Int32Type, Int64Type, FloatType, DoubleType, FLBAType>;

TYPED_TEST_SUITE(EncodingAdHocTyped, EncodingAdHocTypedCases);

TYPED_TEST(EncodingAdHocTyped, PlainArrowDirectPut) {
  for (auto seed : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}) {
    this->Plain(seed);
  }
}

TYPED_TEST(EncodingAdHocTyped, ByteStreamSplitArrowDirectPut) {
  for (auto seed : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}) {
    this->ByteStreamSplit(seed);
  }
}

TEST(DictEncodingAdHoc, ArrowBinaryDirectPut) {
  // Implemented as part of ARROW-3246
  const int64_t size = 50;
  const int64_t min_length = 0;
  const int64_t max_length = 10;
  const double null_probability = 0.1;
  ::arrow::random::RandomArrayGenerator rag(0);
  auto values = rag.String(size, min_length, max_length, null_probability);

  auto owned_encoder = MakeTypedEncoder<ByteArrayType>(Encoding::PLAIN,
                                                       /*use_dictionary=*/true);

  auto encoder = dynamic_cast<DictEncoder<ByteArrayType>*>(owned_encoder.get());

  ASSERT_NO_THROW(encoder->Put(*values));

  std::unique_ptr<ByteArrayDecoder> decoder;
  std::shared_ptr<Buffer> buf, dict_buf;
  int num_values = static_cast<int>(values->length() - values->null_count());
  GetDictDecoder(encoder, num_values, &buf, &dict_buf, nullptr, &decoder);

  typename EncodingTraits<ByteArrayType>::Accumulator acc;
  acc.builder.reset(new ::arrow::StringBuilder);
  ASSERT_EQ(num_values,
            decoder->DecodeArrow(static_cast<int>(values->length()),
                                 static_cast<int>(values->null_count()),
                                 values->null_bitmap_data(), values->offset(), &acc));

  std::shared_ptr<::arrow::Array> result;
  ASSERT_OK(acc.builder->Finish(&result));
  ::arrow::AssertArraysEqual(*values, *result);
}

TYPED_TEST(EncodingAdHocTyped, DictArrowDirectPut) { this->Dict(0); }

TEST(DictEncodingAdHoc, PutDictionaryPutIndices) {
  // Part of ARROW-3246
  auto dict_values =
      ::arrow::ArrayFromJSON(::arrow::binary(), "[\"foo\", \"bar\", \"baz\"]");

  auto CheckIndexType = [&](const std::shared_ptr<::arrow::DataType>& index_ty) {
    auto indices = ::arrow::ArrayFromJSON(index_ty, "[0, 1, 2]");
    auto indices_nulls = ::arrow::ArrayFromJSON(index_ty, "[null, 0, 1, null, 2]");

    auto expected = ::arrow::ArrayFromJSON(::arrow::binary(),
                                           "[\"foo\", \"bar\", \"baz\", null, "
                                           "\"foo\", \"bar\", null, \"baz\"]");

    auto owned_encoder = MakeTypedEncoder<ByteArrayType>(Encoding::PLAIN,
                                                         /*use_dictionary=*/true);
    auto owned_decoder = MakeDictDecoder<ByteArrayType>();

    auto encoder = dynamic_cast<DictEncoder<ByteArrayType>*>(owned_encoder.get());

    ASSERT_NO_THROW(encoder->PutDictionary(*dict_values));

    // Trying to call PutDictionary again throws
    ASSERT_THROW(encoder->PutDictionary(*dict_values), ParquetException);

    ASSERT_NO_THROW(encoder->PutIndices(*indices));
    ASSERT_NO_THROW(encoder->PutIndices(*indices_nulls));

    std::unique_ptr<ByteArrayDecoder> decoder;
    std::shared_ptr<Buffer> buf, dict_buf;
    int num_values = static_cast<int>(expected->length() - expected->null_count());
    GetDictDecoder(encoder, num_values, &buf, &dict_buf, nullptr, &decoder);

    typename EncodingTraits<ByteArrayType>::Accumulator acc;
    acc.builder.reset(new ::arrow::BinaryBuilder);
    ASSERT_EQ(num_values, decoder->DecodeArrow(static_cast<int>(expected->length()),
                                               static_cast<int>(expected->null_count()),
                                               expected->null_bitmap_data(),
                                               expected->offset(), &acc));

    std::shared_ptr<::arrow::Array> result;
    ASSERT_OK(acc.builder->Finish(&result));
    ::arrow::AssertArraysEqual(*expected, *result);
  };

  for (auto ty : ::arrow::all_dictionary_index_types()) {
    CheckIndexType(ty);
  }
}

TYPED_TEST(EncodingAdHocTyped, DictArrowDirectPutIndices) { this->DictPutIndices(); }

class DictEncoding : public TestArrowBuilderDecoding {
 public:
  void SetupEncoderDecoder() override {
    auto node = schema::ByteArray("name");
    descr_ = std::unique_ptr<ColumnDescriptor>(new ColumnDescriptor(node, 0, 0));
    encoder_ = MakeTypedEncoder<ByteArrayType>(Encoding::PLAIN, /*use_dictionary=*/true,
                                               descr_.get());
    if (null_count_ == 0) {
      ASSERT_NO_THROW(encoder_->Put(input_data_.data(), num_values_));
    } else {
      ASSERT_NO_THROW(
          encoder_->PutSpaced(input_data_.data(), num_values_, valid_bits_, 0));
    }
    buffer_ = encoder_->FlushValues();

    auto dict_encoder = dynamic_cast<DictEncoder<ByteArrayType>*>(encoder_.get());
    ASSERT_NE(dict_encoder, nullptr);
    dict_buffer_ =
        AllocateBuffer(default_memory_pool(), dict_encoder->dict_encoded_size());
    dict_encoder->WriteDict(dict_buffer_->mutable_data());

    // Simulate reading the dictionary page followed by a data page
    plain_decoder_ = MakeTypedDecoder<ByteArrayType>(Encoding::PLAIN, descr_.get());
    plain_decoder_->SetData(dict_encoder->num_entries(), dict_buffer_->data(),
                            static_cast<int>(dict_buffer_->size()));

    dict_decoder_ = MakeDictDecoder<ByteArrayType>(descr_.get());
    dict_decoder_->SetDict(plain_decoder_.get());
    dict_decoder_->SetData(num_values_, buffer_->data(),
                           static_cast<int>(buffer_->size()));
    decoder_ = dynamic_cast<ByteArrayDecoder*>(dict_decoder_.get());
  }

 protected:
  std::unique_ptr<ColumnDescriptor> descr_;
  std::shared_ptr<Buffer> dict_buffer_;
};

TEST_F(DictEncoding, CheckDecodeArrowUsingDenseBuilder) {
  this->CheckDecodeArrowUsingDenseBuilder();
}

TEST_F(DictEncoding, CheckDecodeArrowUsingDictBuilder) {
  this->CheckDecodeArrowUsingDictBuilder();
}

TEST_F(DictEncoding, CheckDecodeArrowNonNullDenseBuilder) {
  this->CheckDecodeArrowNonNullUsingDenseBuilder();
}

TEST_F(DictEncoding, CheckDecodeArrowNonNullDictBuilder) {
  this->CheckDecodeArrowNonNullUsingDictBuilder();
}

TEST_F(DictEncoding, CheckDecodeIndicesSpaced) {
  for (auto np : null_probabilities_) {
    InitTestCase(np);
    auto builder = CreateDictBuilder();
    dict_decoder_->InsertDictionary(builder.get());
    int actual_num_values;
    if (null_count_ == 0) {
      actual_num_values = dict_decoder_->DecodeIndices(num_values_, builder.get());
    } else {
      actual_num_values = dict_decoder_->DecodeIndicesSpaced(
          num_values_, null_count_, valid_bits_, 0, builder.get());
    }
    ASSERT_EQ(actual_num_values, num_values_ - null_count_);
    std::shared_ptr<::arrow::Array> actual;
    ASSERT_OK(builder->Finish(&actual));
    ASSERT_ARRAYS_EQUAL(*actual, *expected_dict_);

    // Check that null indices are zero-initialized
    const auto& dict_actual = checked_cast<const ::arrow::DictionaryArray&>(*actual);
    const auto& indices =
        checked_cast<const ::arrow::Int32Array&>(*dict_actual.indices());

    auto raw_values = indices.raw_values();
    for (int64_t i = 0; i < indices.length(); ++i) {
      if (indices.IsNull(i) && raw_values[i] != 0) {
        FAIL() << "Null slot not zero-initialized";
      }
    }
  }
}

TEST_F(DictEncoding, CheckDecodeIndicesNoNulls) {
  InitTestCase(/*null_probability=*/0.0);
  auto builder = CreateDictBuilder();
  dict_decoder_->InsertDictionary(builder.get());
  auto actual_num_values = dict_decoder_->DecodeIndices(num_values_, builder.get());
  CheckDict(actual_num_values, *builder);
}

// ----------------------------------------------------------------------
// BYTE_STREAM_SPLIT encode/decode tests.

template <typename Type>
class TestByteStreamSplitEncoding : public TestEncodingBase<Type> {
 public:
  using c_type = typename Type::c_type;
  static constexpr int TYPE = Type::type_num;

  void CheckRoundtrip() override {
    auto encoder =
        MakeTypedEncoder<Type>(Encoding::BYTE_STREAM_SPLIT, false, descr_.get());
    auto decoder = MakeTypedDecoder<Type>(Encoding::BYTE_STREAM_SPLIT, descr_.get());
    encoder->Put(draws_, num_values_);
    encode_buffer_ = encoder->FlushValues();

    {
      decoder->SetData(num_values_, encode_buffer_->data(),
                       static_cast<int>(encode_buffer_->size()));
      int values_decoded = decoder->Decode(decode_buf_, num_values_);
      ASSERT_EQ(num_values_, values_decoded);
      ASSERT_NO_FATAL_FAILURE(VerifyResults<c_type>(decode_buf_, draws_, num_values_));
    }

    {
      // Try again but with a small step.
      decoder->SetData(num_values_, encode_buffer_->data(),
                       static_cast<int>(encode_buffer_->size()));
      int step = 131;
      int remaining = num_values_;
      for (int i = 0; i < num_values_; i += step) {
        int num_decoded = decoder->Decode(decode_buf_, step);
        ASSERT_EQ(num_decoded, std::min(step, remaining));
        ASSERT_NO_FATAL_FAILURE(
            VerifyResults<c_type>(decode_buf_, &draws_[i], num_decoded));
        remaining -= num_decoded;
      }
    }

    {
      std::vector<uint8_t> valid_bits(::arrow::BitUtil::BytesForBits(num_values_), 0);
      std::vector<c_type> expected_filtered_output;
      const int every_nth = 5;
      expected_filtered_output.reserve((num_values_ + every_nth - 1) / every_nth);
      ::arrow::internal::BitmapWriter writer{valid_bits.data(), 0, num_values_};
      // Set every fifth bit.
      for (int i = 0; i < num_values_; ++i) {
        if (i % every_nth == 0) {
          writer.Set();
          expected_filtered_output.push_back(draws_[i]);
        }
        writer.Next();
      }
      writer.Finish();
      const int expected_size = static_cast<int>(expected_filtered_output.size());
      ASSERT_NO_THROW(encoder->PutSpaced(draws_, num_values_, valid_bits.data(), 0));
      encode_buffer_ = encoder->FlushValues();

      decoder->SetData(expected_size, encode_buffer_->data(),
                       static_cast<int>(encode_buffer_->size()));
      int values_decoded = decoder->Decode(decode_buf_, num_values_);
      ASSERT_EQ(expected_size, values_decoded);
      ASSERT_NO_FATAL_FAILURE(VerifyResults<c_type>(
          decode_buf_, expected_filtered_output.data(), expected_size));
    }
  }

  void CheckDecode();
  void CheckEncode();

 protected:
  USING_BASE_MEMBERS();

  void CheckDecode(const uint8_t* encoded_data, const int64_t encoded_data_size,
                   const c_type* expected_decoded_data, const int num_elements) {
    std::unique_ptr<TypedDecoder<Type>> decoder =
        MakeTypedDecoder<Type>(Encoding::BYTE_STREAM_SPLIT);
    decoder->SetData(num_elements, encoded_data, static_cast<int>(encoded_data_size));
    std::vector<c_type> decoded_data(num_elements);
    int num_decoded_elements = decoder->Decode(decoded_data.data(), num_elements);
    ASSERT_EQ(num_elements, num_decoded_elements);
    for (size_t i = 0U; i < decoded_data.size(); ++i) {
      ASSERT_EQ(expected_decoded_data[i], decoded_data[i]);
    }
    ASSERT_EQ(0, decoder->values_left());
  }

  void CheckEncode(const c_type* data, const int num_elements,
                   const uint8_t* expected_encoded_data,
                   const int64_t encoded_data_size) {
    std::unique_ptr<TypedEncoder<Type>> encoder =
        MakeTypedEncoder<Type>(Encoding::BYTE_STREAM_SPLIT);
    encoder->Put(data, num_elements);
    auto encoded_data = encoder->FlushValues();
    ASSERT_EQ(encoded_data_size, encoded_data->size());
    const uint8_t* encoded_data_raw = encoded_data->data();
    for (int64_t i = 0; i < encoded_data->size(); ++i) {
      ASSERT_EQ(expected_encoded_data[i], encoded_data_raw[i]);
    }
  }
};

template <typename c_type>
static std::vector<c_type> ToLittleEndian(const std::vector<c_type>& input) {
  std::vector<c_type> data(input.size());
  std::transform(input.begin(), input.end(), data.begin(), [](const c_type& value) {
    return ::arrow::BitUtil::ToLittleEndian(value);
  });
  return data;
}

static_assert(sizeof(float) == sizeof(uint32_t),
              "BYTE_STREAM_SPLIT encoding tests assume float / uint32_t type sizes");
static_assert(sizeof(double) == sizeof(uint64_t),
              "BYTE_STREAM_SPLIT encoding tests assume double / uint64_t type sizes");

template <>
void TestByteStreamSplitEncoding<FloatType>::CheckDecode() {
  const uint8_t data[] = {0x11, 0x22, 0x33, 0x44, 0x55, 0x66,
                          0x77, 0x88, 0x99, 0xAA, 0xBB, 0xCC};
  const auto expected_output =
      ToLittleEndian<uint32_t>({0xAA774411U, 0xBB885522U, 0xCC996633U});
  CheckDecode(data, static_cast<int64_t>(sizeof(data)),
              reinterpret_cast<const float*>(expected_output.data()),
              static_cast<int>(sizeof(data) / sizeof(float)));
}

template <>
void TestByteStreamSplitEncoding<DoubleType>::CheckDecode() {
  const uint8_t data[] = {0xDE, 0xC0, 0x37, 0x13, 0x11, 0x22, 0x33, 0x44,
                          0xAA, 0xBB, 0xCC, 0xDD, 0x55, 0x66, 0x77, 0x88};
  const auto expected_output =
      ToLittleEndian<uint64_t>({0x7755CCAA331137DEULL, 0x8866DDBB442213C0ULL});
  CheckDecode(data, static_cast<int64_t>(sizeof(data)),
              reinterpret_cast<const double*>(expected_output.data()),
              static_cast<int>(sizeof(data) / sizeof(double)));
}

template <>
void TestByteStreamSplitEncoding<DoubleType>::CheckEncode() {
  const auto data = ToLittleEndian<uint64_t>(
      {0x4142434445464748ULL, 0x0102030405060708ULL, 0xb1b2b3b4b5b6b7b8ULL});
  const uint8_t expected_output[24] = {
      0x48, 0x08, 0xb8, 0x47, 0x07, 0xb7, 0x46, 0x06, 0xb6, 0x45, 0x05, 0xb5,
      0x44, 0x04, 0xb4, 0x43, 0x03, 0xb3, 0x42, 0x02, 0xb2, 0x41, 0x01, 0xb1,
  };
  CheckEncode(reinterpret_cast<const double*>(data.data()), static_cast<int>(data.size()),
              expected_output, sizeof(expected_output));
}

template <>
void TestByteStreamSplitEncoding<FloatType>::CheckEncode() {
  const auto data = ToLittleEndian<uint32_t>({0xaabbccdd, 0x11223344});
  const uint8_t expected_output[8] = {0xdd, 0x44, 0xcc, 0x33, 0xbb, 0x22, 0xaa, 0x11};
  CheckEncode(reinterpret_cast<const float*>(data.data()), static_cast<int>(data.size()),
              expected_output, sizeof(expected_output));
}

typedef ::testing::Types<FloatType, DoubleType> ByteStreamSplitTypes;
TYPED_TEST_SUITE(TestByteStreamSplitEncoding, ByteStreamSplitTypes);

TYPED_TEST(TestByteStreamSplitEncoding, BasicRoundTrip) {
  for (int values = 0; values < 32; ++values) {
    ASSERT_NO_FATAL_FAILURE(this->Execute(values, 1));
  }

  // We need to test with different sizes to guarantee that the SIMD implementation
  // can handle both inputs with size divisible by 4/8 and sizes which would
  // require a scalar loop for the suffix.
  constexpr size_t kSuffixSize = 7;
  constexpr size_t kAvx2Size = 32;    // sizeof(__m256i) for AVX2
  constexpr size_t kAvx512Size = 64;  // sizeof(__m512i) for AVX512
  constexpr size_t kMultiSimdSize = kAvx512Size * 7;

  // Exercise only one SIMD loop. SSE and AVX2 covered in above loop.
  ASSERT_NO_FATAL_FAILURE(this->Execute(kAvx512Size, 1));
  // Exercise one SIMD loop with suffix. SSE covered in above loop.
  ASSERT_NO_FATAL_FAILURE(this->Execute(kAvx2Size + kSuffixSize, 1));
  ASSERT_NO_FATAL_FAILURE(this->Execute(kAvx512Size + kSuffixSize, 1));
  // Exercise multi SIMD loop.
  ASSERT_NO_FATAL_FAILURE(this->Execute(kMultiSimdSize, 1));
  // Exercise multi SIMD loop with suffix.
  ASSERT_NO_FATAL_FAILURE(this->Execute(kMultiSimdSize + kSuffixSize, 1));
}

TYPED_TEST(TestByteStreamSplitEncoding, RoundTripSingleElement) {
  ASSERT_NO_FATAL_FAILURE(this->Execute(1, 1));
}

TYPED_TEST(TestByteStreamSplitEncoding, CheckOnlyDecode) {
  ASSERT_NO_FATAL_FAILURE(this->CheckDecode());
}

TYPED_TEST(TestByteStreamSplitEncoding, CheckOnlyEncode) {
  ASSERT_NO_FATAL_FAILURE(this->CheckEncode());
}

TEST(ByteStreamSplitEncodeDecode, InvalidDataTypes) {
  // First check encoders.
  ASSERT_THROW(MakeTypedEncoder<Int32Type>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedEncoder<Int64Type>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedEncoder<Int96Type>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedEncoder<BooleanType>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedEncoder<ByteArrayType>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedEncoder<FLBAType>(Encoding::BYTE_STREAM_SPLIT), ParquetException);

  // Then check decoders.
  ASSERT_THROW(MakeTypedDecoder<Int32Type>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedDecoder<Int64Type>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedDecoder<Int96Type>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedDecoder<BooleanType>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedDecoder<ByteArrayType>(Encoding::BYTE_STREAM_SPLIT),
               ParquetException);
  ASSERT_THROW(MakeTypedDecoder<FLBAType>(Encoding::BYTE_STREAM_SPLIT), ParquetException);
}

}  // namespace test
}  // namespace parquet
