GH-50703: [C++][Parquet] Reserve dictionary byte array data

Author: Punisheroot

cpp/src/parquet/decoder.cc

@@ -1024,8 +1024,7 @@ class DictDecoderImpl : public TypedDecoderImpl<Type>, public DictDecoder<Type>
   // memory use in most cases
   std::shared_ptr<ResizableBuffer> byte_array_offsets_;
 
-  // Reusable buffer for decoding dictionary indices to be appended to a
-  // BinaryDictionary32Builder
+  // Reusable buffer for decoding dictionary indices into Arrow builders.
   std::shared_ptr<ResizableBuffer> indices_scratch_space_;
 
   ::arrow::util::RleBitPackedDecoder<int32_t> idx_decoder_;
@@ -1295,12 +1294,80 @@ class DictByteArrayDecoderImpl : public DictDecoderImpl<ByteArrayType> {
                           int64_t valid_bits_offset,
                           typename EncodingTraits<ByteArrayType>::Accumulator* out,
                           int* out_num_values) {
+    const auto* dict_values = dictionary_->data_as<ByteArray>();
+    const int values_to_decode = num_values - null_count;
+
+    switch (out->builder->type()->id()) {
+      case ::arrow::Type::BINARY:
+      case ::arrow::Type::STRING:
+      case ::arrow::Type::LARGE_BINARY:
+      case ::arrow::Type::LARGE_STRING: {
+        if (values_to_decode > 0) {
+          RETURN_NOT_OK(indices_scratch_space_->TypedResize<int32_t>(
+              values_to_decode, /*shrink_to_fit=*/false));
+        }
+        auto* decoded_indices = indices_scratch_space_->mutable_data_as<int32_t>();
+        const int num_indices = idx_decoder_.GetBatch(decoded_indices, values_to_decode);
+        if (ARROW_PREDICT_FALSE(num_indices != values_to_decode)) {
+          return Status::Invalid("Invalid number of indices: ", num_indices);
+        }
+
+        int64_t data_length = 0;
+        for (int i = 0; i < values_to_decode; ++i) {
+          const auto index = decoded_indices[i];
+          RETURN_NOT_OK(IndexInBounds(index));
+          if (ARROW_PREDICT_FALSE(AddWithOverflow(
+                  data_length, static_cast<int64_t>(dict_values[index].len),
+                  &data_length))) {
+            return Status::Invalid(
+                "excess expansion while decoding dictionary-encoded BYTE_ARRAY");
+          }
+        }
+
+        auto append_predecoded = [&](auto* helper) {
+          int values_decoded = 0;
+          int pos_indices = 0;
+          int64_t remaining_data_length = data_length;
+
+          RETURN_NOT_OK(VisitBitRuns(
+              valid_bits, valid_bits_offset, num_values,
+              [&](int64_t position, int64_t length, bool valid) {
+                if (valid) {
+                  for (int64_t i = 0; i < length; ++i) {
+                    const auto& val = dict_values[decoded_indices[pos_indices++]];
+                    RETURN_NOT_OK(helper->AppendValue(
+                        val.ptr, static_cast<int32_t>(val.len), remaining_data_length));
+                    remaining_data_length -= val.len;
+                  }
+                  values_decoded += static_cast<int>(length);
+                } else {
+                  for (int64_t i = 0; i < length; ++i) {
+                    helper->UnsafeAppendNull();
+                  }
+                }
+                return Status::OK();
+              }));
+          DCHECK_EQ(pos_indices, values_to_decode);
+          DCHECK_EQ(remaining_data_length, 0);
+          *out_num_values = values_decoded;
+          return Status::OK();
+        };
+
+        return DispatchArrowBinaryHelper<ByteArrayType>(out, num_values, data_length,
+                                                        append_predecoded);
+      }
+      default:
+        // Binary-view builders don't benefit from reserving the dictionary values'
+        // total byte length, since short values are stored inline. Keep their
+        // existing bounded streaming decode path. Unsupported builder types are
+        // rejected by DispatchArrowBinaryHelper below before decoding any indices.
+        break;
+    }
+
     constexpr int32_t kBufferSize = 1024;
     int32_t indices[kBufferSize];
 
-    auto visit_binary_helper = [&](auto* helper) {
-      const auto* dict_values = dictionary_->data_as<ByteArray>();
-      const int values_to_decode = num_values - null_count;
+    auto append_streaming = [&](auto* helper) {
       int values_decoded = 0;
       int num_indices = 0;
       int pos_indices = 0;
@@ -1309,7 +1376,7 @@ class DictByteArrayDecoderImpl : public DictDecoderImpl<ByteArrayType> {
         if (valid) {
           while (length > 0) {
             if (num_indices == pos_indices) {
-              // Refill indices buffer
+              // Refill the bounded indices buffer for binary-view builders.
               const auto max_batch_size =
                   std::min<int32_t>(kBufferSize, values_to_decode - values_decoded);
               num_indices = idx_decoder_.GetBatch(indices, max_batch_size);
@@ -1341,11 +1408,9 @@ class DictByteArrayDecoderImpl : public DictDecoderImpl<ByteArrayType> {
       *out_num_values = values_decoded;
       return Status::OK();
     };
-    // The `len_` in the ByteArrayDictDecoder is the total length of the
-    // RLE/Bit-pack encoded data size, so, we cannot use `len_` to reserve
-    // space for binary data.
+
     return DispatchArrowBinaryHelper<ByteArrayType>(
-        out, num_values, /*estimated_data_length=*/{}, visit_binary_helper);
+        out, num_values, /*estimated_data_length=*/{}, append_streaming);
   }
 
   template <typename BuilderType>

cpp/src/parquet/encoding_benchmark.cc

@@ -1300,14 +1300,15 @@ class BenchmarkDecodeArrowByteArray : public BenchmarkDecodeArrowBase<ByteArrayT
   }
 
   void InitDataInputs() final {
-    // Generate a random string dictionary without any nulls so that this dataset can
-    // be used for benchmarking the DecodeArrowNonNull API
+    // The default dataset has no nulls so that it can also be used for benchmarking
+    // the DecodeArrowNonNull API. DecodeArrowWithNullDenseBenchmark regenerates it
+    // with the requested null probability.
     constexpr int repeat_factor = 8;
     constexpr int64_t min_length = 2;
     constexpr int64_t max_length = 10;
     ::arrow::random::RandomArrayGenerator rag(0);
     input_array_ = rag.StringWithRepeats(num_values_, num_values_ / repeat_factor,
-                                         min_length, max_length, /*null_probability=*/0);
+                                         min_length, max_length, null_probability_);
     valid_bits_ = input_array_->null_bitmap_data();
     total_size_ = input_array_->data()->buffers[2]->size();
 
@@ -1318,7 +1319,23 @@ class BenchmarkDecodeArrowByteArray : public BenchmarkDecodeArrowBase<ByteArrayT
     }
   }
 
+  void DecodeArrowWithNullDenseBenchmark(benchmark::State& state) {
+    null_probability_ = static_cast<double>(state.range(1)) / 10000;
+    InitDataInputs();
+    DoEncodeArrow();
+
+    for (auto _ : state) {
+      auto decoder = InitializeDecoder();
+      auto acc = CreateAccumulator();
+      decoder->DecodeArrow(num_values_, static_cast<int>(input_array_->null_count()),
+                           valid_bits_, 0, &acc);
+    }
+    state.SetBytesProcessed(state.iterations() * total_size_);
+    state.SetItemsProcessed(state.iterations() * num_values_);
+  }
+
  protected:
+  double null_probability_{0.0};
   std::vector<ByteArray> values_;
 };
 
@@ -1478,6 +1495,14 @@ class BM_ArrowBinaryViewDict : public BM_ArrowBinaryDict {
   }
 };
 
+static void ByteArrayWithNullCustomArguments(benchmark::internal::Benchmark* b) {
+  b->ArgsProduct({
+                     benchmark::CreateRange(MIN_RANGE, MAX_RANGE, /*multi=*/4),
+                     {1000, 5000},
+                 })
+      ->ArgNames({"num_values", "null_in_ten_thousand"});
+}
+
 BENCHMARK_DEFINE_F(BM_ArrowBinaryDict, EncodeArrow)
 (benchmark::State& state) { EncodeArrowBenchmark(state); }
 BENCHMARK_REGISTER_F(BM_ArrowBinaryDict, EncodeArrow)->Range(1 << 18, 1 << 20);
@@ -1507,6 +1532,11 @@ BENCHMARK_DEFINE_F(BM_ArrowBinaryDict, DecodeArrow_Dense)(benchmark::State& stat
 }
 BENCHMARK_REGISTER_F(BM_ArrowBinaryDict, DecodeArrow_Dense)->Range(MIN_RANGE, MAX_RANGE);
 
+BENCHMARK_DEFINE_F(BM_ArrowBinaryDict, DecodeArrowWithNull_Dense)
+(benchmark::State& state) { DecodeArrowWithNullDenseBenchmark(state); }
+BENCHMARK_REGISTER_F(BM_ArrowBinaryDict, DecodeArrowWithNull_Dense)
+    ->Apply(ByteArrayWithNullCustomArguments);
+
 BENCHMARK_DEFINE_F(BM_ArrowBinaryDict, DecodeArrowNonNull_Dense)
 (benchmark::State& state) { DecodeArrowNonNullDenseBenchmark(state); }
 BENCHMARK_REGISTER_F(BM_ArrowBinaryDict, DecodeArrowNonNull_Dense)
@@ -1527,6 +1557,11 @@ BENCHMARK_DEFINE_F(BM_ArrowBinaryViewDict, DecodeArrow_Dense)(benchmark::State&
 BENCHMARK_REGISTER_F(BM_ArrowBinaryViewDict, DecodeArrow_Dense)
     ->Range(MIN_RANGE, MAX_RANGE);
 
+BENCHMARK_DEFINE_F(BM_ArrowBinaryViewDict, DecodeArrowWithNull_Dense)
+(benchmark::State& state) { DecodeArrowWithNullDenseBenchmark(state); }
+BENCHMARK_REGISTER_F(BM_ArrowBinaryViewDict, DecodeArrowWithNull_Dense)
+    ->Apply(ByteArrayWithNullCustomArguments);
+
 BENCHMARK_DEFINE_F(BM_ArrowBinaryViewDict, DecodeArrowNonNull_Dense)
 (benchmark::State& state) { DecodeArrowNonNullDenseBenchmark(state); }
 BENCHMARK_REGISTER_F(BM_ArrowBinaryViewDict, DecodeArrowNonNull_Dense)

cpp/src/parquet/encoding_test.cc

@@ -41,6 +41,7 @@
 #include "arrow/util/bitmap_writer.h"
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/endian.h"
+#include "arrow/util/rle_encoding_internal.h"
 #include "arrow/util/string.h"
 #include "parquet/encoding.h"
 #include "parquet/platform.h"
@@ -1265,6 +1266,54 @@ TEST(DictEncodingAdHoc, ArrowBinaryDirectPut) {
   ::arrow::AssertArraysEqual(*values, *result);
 }
 
+TEST(DictEncodingAdHoc, DenseDecodeRejectsInvalidOrTruncatedIndices) {
+  auto dictionary =
+      ::arrow::ArrayFromJSON(::arrow::binary(), R"(["a", "bb", "ccc"])");
+  auto owned_encoder = MakeTypedEncoder<ByteArrayType>(
+      Encoding::PLAIN, /*use_dictionary=*/true);
+  auto* encoder = dynamic_cast<DictEncoder<ByteArrayType>*>(owned_encoder.get());
+  ASSERT_NE(encoder, nullptr);
+  ASSERT_NO_THROW(encoder->PutDictionary(*dictionary));
+
+  auto dictionary_buffer =
+      AllocateBuffer(default_memory_pool(), encoder->dict_encoded_size());
+  encoder->WriteDict(dictionary_buffer->mutable_data());
+
+  auto dictionary_decoder =
+      MakeTypedDecoder<ByteArrayType>(Encoding::PLAIN, /*descr=*/nullptr);
+  dictionary_decoder->SetData(encoder->num_entries(), dictionary_buffer->data(),
+                              static_cast<int>(dictionary_buffer->size()));
+
+  auto decoder = MakeDictDecoder<ByteArrayType>();
+  decoder->SetDict(dictionary_decoder.get());
+
+  auto ExpectDecodeFailure = [&](const uint8_t* data, int size) {
+    decoder->SetData(/*num_values=*/1, data, size);
+    typename EncodingTraits<ByteArrayType>::Accumulator acc;
+    acc.builder = std::make_unique<::arrow::BinaryBuilder>();
+    ASSERT_THROW(
+        decoder->DecodeArrow(/*num_values=*/1, /*null_count=*/0,
+                             /*valid_bits=*/nullptr, /*valid_bits_offset=*/0, &acc),
+        ParquetException);
+  };
+
+  constexpr int kBitWidth = 2;
+  std::vector<uint8_t> invalid_index_data(
+      1 + ::arrow::util::RleBitPackedEncoder::MaxBufferSize(
+              kBitWidth, /*num_values=*/1) +
+      ::arrow::util::RleBitPackedEncoder::MinBufferSize(kBitWidth));
+  invalid_index_data[0] = kBitWidth;
+  ::arrow::util::RleBitPackedEncoder index_encoder(
+      invalid_index_data.data() + 1,
+      static_cast<int>(invalid_index_data.size() - 1), kBitWidth);
+  ASSERT_TRUE(index_encoder.Put(/*value=*/3));
+  const int invalid_index_size = 1 + index_encoder.Flush();
+  ExpectDecodeFailure(invalid_index_data.data(), invalid_index_size);
+
+  const uint8_t truncated_index_data[] = {kBitWidth};
+  ExpectDecodeFailure(truncated_index_data, sizeof(truncated_index_data));
+}
+
 TEST(DictEncodingAdHoc, PutDictionaryPutIndices) {
   // Part of ARROW-3246
   auto dict_values =