GH-32863: [C++][Parquet] Add DELTA_BYTE_ARRAY encoder to Parquet writer

cpp/src/parquet/encoding.h

@@ -141,13 +141,13 @@ struct EncodingTraits<ByteArrayType> {
  using Encoder = ByteArrayEncoder;
  using Decoder = ByteArrayDecoder;
 
+ using ArrowType = ::arrow::BinaryType;
  /// \brief Internal helper class for decoding BYTE_ARRAY data where we can
  /// overflow the capacity of a single arrow::BinaryArray
  struct Accumulator {
  std::unique_ptr<::arrow::BinaryBuilder> builder;
  std::vector<std::shared_ptr<::arrow::Array>> chunks;
  };
- using ArrowType = ::arrow::BinaryType;
  using DictAccumulator = ::arrow::Dictionary32Builder<::arrow::BinaryType>;
 };
 

cpp/src/parquet/encoding_test.cc

@@ -36,6 +36,7 @@
 #include "arrow/util/bitmap_writer.h"
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/endian.h"
+#include "arrow/util/span.h"
 #include "arrow/util/string.h"
 #include "parquet/encoding.h"
 #include "parquet/platform.h"
@@ -181,7 +182,7 @@ class TestEncodingBase : public ::testing::Test {
 
  void TearDown() {}
 
- void InitData(int nvalues, int repeats) {
+virtual 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));
@@ -1705,11 +1706,13 @@ class TestDeltaLengthByteArrayEncoding : public TestEncodingBase<Type> {
  using c_type = typename Type::c_type;
  static constexpr int TYPE = Type::type_num;
 
+ virtual Encoding::type GetEncoding() { return Encoding::DELTA_LENGTH_BYTE_ARRAY; }
+
  virtual void CheckRoundtrip() {
- auto encoder = MakeTypedEncoder<Type>(Encoding::DELTA_LENGTH_BYTE_ARRAY,
+ auto encoding = GetEncoding();
+ auto encoder = MakeTypedEncoder<Type>(encoding,
  /*use_dictionary=*/false, descr_.get());
- auto decoder =
- MakeTypedDecoder<Type>(Encoding::DELTA_LENGTH_BYTE_ARRAY, descr_.get());
+ auto decoder = MakeTypedDecoder<Type>(encoding, descr_.get());
 
  encoder->Put(draws_, num_values_);
  encode_buffer_ = encoder->FlushValues();
@@ -1722,10 +1725,10 @@ class TestDeltaLengthByteArrayEncoding : public TestEncodingBase<Type> {
  }
 
  void CheckRoundtripSpaced(const uint8_t* valid_bits, int64_t valid_bits_offset) {
- auto encoder = MakeTypedEncoder<Type>(Encoding::DELTA_LENGTH_BYTE_ARRAY,
+ auto encoding = GetEncoding();
+ auto encoder = MakeTypedEncoder<Type>(encoding,
  /*use_dictionary=*/false, descr_.get());
- auto decoder =
- MakeTypedDecoder<Type>(Encoding::DELTA_LENGTH_BYTE_ARRAY, descr_.get());
+ auto decoder = MakeTypedDecoder<Type>(encoding, descr_.get());
  int null_count = 0;
  for (auto i = 0; i < num_values_; i++) {
  if (!bit_util::GetBit(valid_bits, valid_bits_offset + i)) {
@@ -1771,6 +1774,19 @@ std::shared_ptr<Buffer> DeltaEncode(std::vector<int32_t> lengths) {
  return encoder->FlushValues();
 }
 
+std::shared_ptr<Buffer> DeltaEncode(::arrow::util::span<const int32_t> lengths) {
+ auto encoder = MakeTypedEncoder<Int32Type>(Encoding::DELTA_BINARY_PACKED);
+ encoder->Put(lengths.data(), static_cast<int>(lengths.size()));
+ return encoder->FlushValues();
+}
+
+std::shared_ptr<Buffer> DeltaEncode(std::shared_ptr<::arrow::Array>& lengths) {
+ auto data = ::arrow::internal::checked_pointer_cast<const ::arrow::Int32Array>(lengths);
+ auto span = ::arrow::util::span<const int32_t>{data->raw_values(),
+ static_cast<size_t>(lengths->length())};
+ return DeltaEncode(span);
+}
+
 TEST(TestDeltaLengthByteArrayEncoding, AdHocRoundTrip) {
  const std::shared_ptr<::arrow::Array> cases[] = {
  ::arrow::ArrayFromJSON(::arrow::utf8(), R"([])"),
@@ -1780,10 +1796,10 @@ TEST(TestDeltaLengthByteArrayEncoding, AdHocRoundTrip) {
  };
 
  std::string expected_encoded_vals[] = {
- DeltaEncode({})->ToString(),
- DeltaEncode({3, 2, 0})->ToString() + "abcde",
- DeltaEncode({0, 0, 0})->ToString(),
- DeltaEncode({0, 3})->ToString() + "xyz",
+ DeltaEncode(std::vector<int>({}))->ToString(),
+ DeltaEncode(std::vector<int>({3, 2, 0}))->ToString() + "abcde",
+ DeltaEncode(std::vector<int>({0, 0, 0}))->ToString(),
+ DeltaEncode(std::vector<int>({0, 3}))->ToString() + "xyz",
  };
 
  auto encoder = MakeTypedEncoder<ByteArrayType>(Encoding::DELTA_LENGTH_BYTE_ARRAY,
@@ -1894,7 +1910,6 @@ TEST(DeltaLengthByteArrayEncodingAdHoc, ArrowBinaryDirectPut) {
  ASSERT_EQ(values->length(), result->length());
  ASSERT_OK(result->ValidateFull());
 
- auto upcast_result = CastBinaryTypesHelper(result, values->type());
  ::arrow::AssertArraysEqual(*values, *result);
  };
 
@@ -1977,4 +1992,263 @@ TEST(DeltaLengthByteArrayEncodingAdHoc, ArrowDirectPut) {
  CheckDecode(encoded, ::arrow::ArrayFromJSON(::arrow::large_binary(), values));
 }
 
+// ----------------------------------------------------------------------
+// DELTA_BYTE_ARRAY encode/decode tests.
+
+template <typename Type>
+class TestDeltaByteArrayEncoding : public TestDeltaLengthByteArrayEncoding<Type> {
+ public:
+ using c_type = typename Type::c_type;
+ static constexpr int TYPE = Type::type_num;
+ static constexpr double prefixed_probability = 0.5;
+
+ void InitData(int nvalues, int repeats) override {
+ 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());
+ GeneratePrefixedData<c_type>(nvalues, draws_, &data_buffer_, prefixed_probability);
+
+ // add some repeated values
+ for (int j = 1; j < repeats; ++j) {
+ for (int i = 0; i < nvalues; ++i) {
+ draws_[nvalues * j + i] = draws_[i];
+ }
+ }
+ }
+
+ Encoding::type GetEncoding() override { return Encoding::DELTA_BYTE_ARRAY; }
+
+ protected:
+ USING_BASE_MEMBERS();
+ std::vector<uint8_t> input_bytes_;
+ std::vector<uint8_t> output_bytes_;
+};
+
+using TestDeltaByteArrayEncodingTypes = ::testing::Types<ByteArrayType, FLBAType>;
+TYPED_TEST_SUITE(TestDeltaByteArrayEncoding, TestDeltaByteArrayEncodingTypes);
+
+TYPED_TEST(TestDeltaByteArrayEncoding, BasicRoundTrip) {
+ ASSERT_NO_FATAL_FAILURE(this->Execute(0, /*repeats=*/0));
+ ASSERT_NO_FATAL_FAILURE(this->Execute(250, 5));
+ ASSERT_NO_FATAL_FAILURE(this->Execute(2000, 1));
+ ASSERT_NO_FATAL_FAILURE(this->ExecuteSpaced(
+ /*nvalues*/ 1234, /*repeats*/ 1, /*valid_bits_offset*/ 64, /*null_probability*/
+ 0));
+ ASSERT_NO_FATAL_FAILURE(this->ExecuteSpaced(
+ /*nvalues*/ 1234, /*repeats*/ 10, /*valid_bits_offset*/ 64,
+ /*null_probability*/ 0.5));
+}
+
+template <typename Type>
+class TestDeltaByteArrayEncodingDirectPut : public TestEncodingBase<Type> {
+ using ArrowType = typename EncodingTraits<Type>::ArrowType;
+ using Accumulator = typename EncodingTraits<Type>::Accumulator;
+ using BuilderType = typename ::arrow::TypeTraits<ArrowType>::BuilderType;
+
+ public:
+ std::unique_ptr<TypedEncoder<Type>> encoder =
+ MakeTypedEncoder<Type>(Encoding::DELTA_BYTE_ARRAY);
+ std::unique_ptr<TypedDecoder<Type>> decoder =
+ MakeTypedDecoder<Type>(Encoding::DELTA_BYTE_ARRAY);
+
+ void CheckDirectPut(std::shared_ptr<::arrow::Array> array);
+
+ void CheckRoundtrip() override;
+
+ protected:
+ USING_BASE_MEMBERS();
+};
+
+template <>
+void TestDeltaByteArrayEncodingDirectPut<ByteArrayType>::CheckDirectPut(
+ std::shared_ptr<::arrow::Array> array) {
+ ASSERT_NO_THROW(encoder->Put(*array));
+ auto buf = encoder->FlushValues();
+
+ int num_values = static_cast<int>(array->length() - array->null_count());
+ decoder->SetData(num_values, buf->data(), static_cast<int>(buf->size()));
+
+ Accumulator acc;
+ acc.builder = std::make_unique<BuilderType>(array->type(), default_memory_pool());
+
+ ASSERT_EQ(num_values,
+ decoder->DecodeArrow(static_cast<int>(array->length()),
+ static_cast<int>(array->null_count()),
+ array->null_bitmap_data(), array->offset(), &acc));
+
+ ASSERT_EQ(acc.chunks.size(), 0) << "Accumulator shouldn't have overflowed chunks";
+ ASSERT_OK_AND_ASSIGN(auto result, acc.builder->Finish());
+ ASSERT_EQ(array->length(), result->length());
+ ASSERT_OK(result->ValidateFull());
+
+ ::arrow::AssertArraysEqual(*array, *result);
+}
+
+template <>
+void TestDeltaByteArrayEncodingDirectPut<FLBAType>::CheckDirectPut(
+ std::shared_ptr<::arrow::Array> array) {
+ ASSERT_NO_THROW(encoder->Put(*array));
+ auto buf = encoder->FlushValues();
+
+ int num_values = static_cast<int>(array->length() - array->null_count());
+ decoder->SetData(num_values, buf->data(), static_cast<int>(buf->size()));
+
+ Accumulator acc(array->type(), default_memory_pool());
+
+ ASSERT_EQ(num_values,
+ decoder->DecodeArrow(static_cast<int>(array->length()),
+ static_cast<int>(array->null_count()),
+ array->null_bitmap_data(), array->offset(), &acc));
+
+ ASSERT_OK_AND_ASSIGN(auto result, acc.Finish());
+ ASSERT_EQ(array->length(), result->length());
+ ASSERT_OK(result->ValidateFull());
+
+ ::arrow::AssertArraysEqual(*array, *result);
+}
+
+template <>
+void TestDeltaByteArrayEncodingDirectPut<ByteArrayType>::CheckRoundtrip() {
+ constexpr int64_t kSize = 500;
+ constexpr int32_t kMinLength = 0;
+ constexpr int32_t kMaxLength = 10;
+ constexpr int32_t kNumUnique = 10;
+ constexpr double kNullProbability = 0.25;
+ constexpr int kSeed = 42;
+ ::arrow::random::RandomArrayGenerator rag{kSeed};
+ std::shared_ptr<::arrow::Array> values = rag.BinaryWithRepeats(
+ /*size=*/1, /*unique=*/1, kMinLength, kMaxLength, kNullProbability);
+ CheckDirectPut(values);
+
+ for (int i = 0; i < 10; ++i) {
+ values = rag.BinaryWithRepeats(kSize, kNumUnique, kMinLength, kMaxLength,
+ kNullProbability);
+ CheckDirectPut(values);
+ }
+}
+
+template <>
+void TestDeltaByteArrayEncodingDirectPut<FLBAType>::CheckRoundtrip() {
+ constexpr int64_t kSize = 50;
+ constexpr int kSeed = 42;
+ constexpr int kByteWidth = 4;
+ ::arrow::random::RandomArrayGenerator rag{kSeed};
+ std::shared_ptr<::arrow::Array> values =
+ rag.FixedSizeBinary(/*size=*/0, /*byte_width=*/kByteWidth);
+ CheckDirectPut(values);
+
+ for (auto seed : {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}) {
+ values = rag.FixedSizeBinary(kSize + seed, kByteWidth);
+ CheckDirectPut(values);
+ }
+}
+
+TYPED_TEST_SUITE(TestDeltaByteArrayEncodingDirectPut, TestDeltaByteArrayEncodingTypes);
+
+TYPED_TEST(TestDeltaByteArrayEncodingDirectPut, DirectPut) {
+ ASSERT_NO_FATAL_FAILURE(this->CheckRoundtrip());
+}
+
+TEST(DeltaByteArrayEncodingAdHoc, ArrowDirectPut) {
+ auto CheckEncode = [](const std::shared_ptr<::arrow::Array>& values,
+ const std::shared_ptr<Buffer>& encoded) {
+ auto encoder = MakeTypedEncoder<ByteArrayType>(Encoding::DELTA_BYTE_ARRAY);
+ ASSERT_NO_THROW(encoder->Put(*values));
+ auto buf = encoder->FlushValues();
+ ASSERT_TRUE(encoded->Equals(*buf));
+ };
+
+ auto CheckDecode = [](std::shared_ptr<Buffer> buf,
+ std::shared_ptr<::arrow::Array> values) {
+ int num_values = static_cast<int>(values->length());
+ auto decoder = MakeTypedDecoder<ByteArrayType>(Encoding::DELTA_BYTE_ARRAY);
+ decoder->SetData(num_values, buf->data(), static_cast<int>(buf->size()));
+
+ typename EncodingTraits<ByteArrayType>::Accumulator acc;
+ if (::arrow::is_string(values->type()->id())) {
+ acc.builder = std::make_unique<::arrow::StringBuilder>();
+ } else {
+ acc.builder = std::make_unique<::arrow::BinaryBuilder>();
+ }
+
+ 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(num_values, result->length());
+ ASSERT_OK(result->ValidateFull());
+
+ auto upcast_result = CastBinaryTypesHelper(result, values->type());
+ ::arrow::AssertArraysEqual(*values, *upcast_result);
+ };
+
+ auto CheckEncodeDecode = [&](std::string_view values,
+ std::shared_ptr<::arrow::Array> prefix_lengths,
+ std::shared_ptr<::arrow::Array> suffix_lengths,
+ std::string_view suffix_data) {
+ auto encoded = ::arrow::ConcatenateBuffers({DeltaEncode(prefix_lengths),
+ DeltaEncode(suffix_lengths),
+ std::make_shared<Buffer>(suffix_data)})
+ .ValueOrDie();
+
+ CheckEncode(::arrow::ArrayFromJSON(::arrow::utf8(), values), encoded);
+ CheckEncode(::arrow::ArrayFromJSON(::arrow::large_utf8(), values), encoded);
+ CheckEncode(::arrow::ArrayFromJSON(::arrow::binary(), values), encoded);
+ CheckEncode(::arrow::ArrayFromJSON(::arrow::large_binary(), values), encoded);
+
+ CheckDecode(encoded, ::arrow::ArrayFromJSON(::arrow::utf8(), values));
+ CheckDecode(encoded, ::arrow::ArrayFromJSON(::arrow::large_utf8(), values));
+ CheckDecode(encoded, ::arrow::ArrayFromJSON(::arrow::binary(), values));
+ CheckDecode(encoded, ::arrow::ArrayFromJSON(::arrow::large_binary(), values));
+ };
+
+ {
+ auto values = R"(["axis", "axle", "babble", "babyhood"])";
+ auto prefix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([0, 2, 0, 3])");
+ auto suffix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([4, 2, 6, 5])");
+
+ constexpr std::string_view suffix_data = "axislebabbleyhood";
+ CheckEncodeDecode(values, prefix_lengths, suffix_lengths, suffix_data);
+ }
+
+ {
+ auto values = R"(["axis", "axis", "axis", "axis"])";
+ auto prefix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([0, 4, 4, 4])");
+ auto suffix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([4, 0, 0, 0])");
+
+ constexpr std::string_view suffix_data = "axis";
+ CheckEncodeDecode(values, prefix_lengths, suffix_lengths, suffix_data);
+ }
+
+ {
+ auto values = R"(["axisba", "axis", "axis", "axis"])";
+ auto prefix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([0, 4, 4, 4])");
+ auto suffix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([6, 0, 0, 0])");
+
+ constexpr std::string_view suffix_data = "axisba";
+ CheckEncodeDecode(values, prefix_lengths, suffix_lengths, suffix_data);
+ }
+
+ {
+ auto values = R"(["baaxis", "axis", "axis", "axis"])";
+ auto prefix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([0, 0, 4, 4])");
+ auto suffix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([6, 4, 0, 0])");
+
+ constexpr std::string_view suffix_data = "baaxisaxis";
+ CheckEncodeDecode(values, prefix_lengths, suffix_lengths, suffix_data);
+ }
+
+ {
+ auto values = R"(["καλημέρα", "καμηλιέρη", "καμηλιέρη", "καλημέρα"])";
+ auto prefix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([0, 5, 18, 5])");
+ auto suffix_lengths = ::arrow::ArrayFromJSON(::arrow::int32(), R"([16, 13, 0, 11])");
+ const std::string suffix_data = "καλημέρα\xbcηλιέρη\xbbημέρα";
+ CheckEncodeDecode(values, prefix_lengths, suffix_lengths, suffix_data);
+ }
+}
 } // namespace parquet::test