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

cpp/src/parquet/decoder.cc

@@ -97,6 +97,20 @@ struct ArrowBinaryHelper<ByteArrayType, ::arrow::BinaryType> {
     return Status::OK();
   }
 
+  // Keep the common per-value path small when callers already know the remaining size.
+  ARROW_FORCE_INLINE Status AppendValue(const uint8_t* data, int32_t length,
+                                        int64_t estimated_remaining_data_length) {
+    DCHECK_GT(entries_remaining_, 0);
+
+    if (ARROW_PREDICT_FALSE(!CanFit(length))) {
+      return AppendValueWithKnownSizeSlow(data, length, estimated_remaining_data_length);
+    }
+    chunk_space_remaining_ -= length;
+    --entries_remaining_;
+    builder_->UnsafeAppend(data, length);
+    return Status::OK();
+  }
+
   // If a new chunk is created and estimated_remaining_data_length is provided,
   // it will also reserve the estimated data length for this chunk.
   Status AppendValue(const uint8_t* data, int32_t length,
@@ -133,6 +147,9 @@ struct ArrowBinaryHelper<ByteArrayType, ::arrow::BinaryType> {
   }
 
  private:
+  ARROW_NOINLINE Status AppendValueWithKnownSizeSlow(
+      const uint8_t* data, int32_t length, int64_t estimated_remaining_data_length);
+
   Status PushChunk() {
     ARROW_ASSIGN_OR_RAISE(auto chunk, acc_->builder->Finish());
     acc_->chunks.push_back(std::move(chunk));
@@ -1024,8 +1041,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_;
@@ -1257,6 +1273,11 @@ void DictDecoderImpl<ByteArrayType>::InsertDictionary(::arrow::ArrayBuilder* bui
   PARQUET_THROW_NOT_OK(binary_builder->InsertMemoValues(*arr));
 }
 
+// Keep defensive bitmap validation out of the hot dictionary decode function.
+ARROW_NOINLINE void ValidateBinaryValidityBitmap(int num_values, int null_count,
+                                                 const uint8_t* valid_bits,
+                                                 int64_t valid_bits_offset);
+
 class DictByteArrayDecoderImpl : public DictDecoderImpl<ByteArrayType> {
  public:
   using BASE = DictDecoderImpl<ByteArrayType>;
@@ -1284,6 +1305,17 @@ class DictByteArrayDecoderImpl : public DictDecoderImpl<ByteArrayType> {
                                             /*valid_bits=*/nullptr, valid_bits_offset,
                                             out, &result));
     } else {
+      switch (out->builder->type()->id()) {
+        case ::arrow::Type::BINARY:
+        case ::arrow::Type::STRING:
+        case ::arrow::Type::LARGE_BINARY:
+        case ::arrow::Type::LARGE_STRING:
+          ValidateBinaryValidityBitmap(num_values, null_count, valid_bits,
+                                       valid_bits_offset);
+          break;
+        default:
+          break;
+      }
       PARQUET_THROW_NOT_OK(DecodeArrowDense(num_values, null_count, valid_bits,
                                             valid_bits_offset, out, &result));
     }
@@ -1295,12 +1327,82 @@ 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 or truncated dictionary index stream: expected ", values_to_decode,
+              " indices but decoded ", 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 +1411,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 +1443,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>
@@ -2372,6 +2472,33 @@ class ByteStreamSplitDecoder<FLBAType> : public ByteStreamSplitDecoderBase<FLBAT
   }
 };
 
+// Keep chunk rollover out of the hot decoder functions that instantiate this helper.
+Status
+ArrowBinaryHelper<ByteArrayType, ::arrow::BinaryType>::AppendValueWithKnownSizeSlow(
+    const uint8_t* data, int32_t length, int64_t estimated_remaining_data_length) {
+  RETURN_NOT_OK(PushChunk());
+  // Validate that the first value fits in an empty chunk before UnsafeAppend below.
+  RETURN_NOT_OK(builder_->ReserveData(length));
+  RETURN_NOT_OK(ReserveInitialChunkData(estimated_remaining_data_length));
+  chunk_space_remaining_ -= length;
+  --entries_remaining_;
+  builder_->UnsafeAppend(data, length);
+  return Status::OK();
+}
+
+void ValidateBinaryValidityBitmap(int num_values, int null_count,
+                                  const uint8_t* valid_bits, int64_t valid_bits_offset) {
+  const int64_t actual_non_null =
+      valid_bits == nullptr
+          ? num_values
+          : ::arrow::internal::CountSetBits(valid_bits, valid_bits_offset, num_values);
+  if (ARROW_PREDICT_FALSE(actual_non_null != num_values - null_count)) {
+    throw ParquetException(
+        "Invalid validity bitmap: null_count does not match the number of non-null "
+        "values");
+  }
+}
+
 }  // namespace
 
 // ----------------------------------------------------------------------

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,82 @@ TEST(DictEncodingAdHoc, ArrowBinaryDirectPut) {
   ::arrow::AssertArraysEqual(*values, *result);
 }
 
+TEST(DictEncodingAdHoc, DenseDecodeRejectsInvalidInput) {
+  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));
+
+  auto valid_indices = ::arrow::ArrayFromJSON(::arrow::int32(), R"([0, 1])");
+  ASSERT_NO_THROW(encoder->PutIndices(*valid_indices));
+  auto valid_index_data = encoder->FlushValues();
+
+  auto ExpectValidityBitmapFailure = [&](const char* case_name, int num_values,
+                                         int null_count, uint8_t valid_bits) {
+    SCOPED_TRACE(case_name);
+    for (const auto& type : {::arrow::binary(), ::arrow::utf8(), ::arrow::large_binary(),
+                             ::arrow::large_utf8()}) {
+      SCOPED_TRACE(type->ToString());
+      decoder->SetData(num_values, valid_index_data->data(),
+                       static_cast<int>(valid_index_data->size()));
+      typename EncodingTraits<ByteArrayType>::Accumulator acc;
+      ASSERT_OK_AND_ASSIGN(acc.builder, ::arrow::MakeBuilder(type));
+      ASSERT_THROW(decoder->DecodeArrow(num_values, null_count, &valid_bits,
+                                        /*valid_bits_offset=*/0, &acc),
+                   ParquetException);
+    }
+  };
+
+  // More non-null slots than promised by null_count must not read past the
+  // predecoded indices buffer.
+  ExpectValidityBitmapFailure("more non-null slots", /*num_values=*/2,
+                              /*null_count=*/1,
+                              /*valid_bits=*/0b00000011);
+  // Fewer non-null slots must not silently leave predecoded indices unconsumed.
+  ExpectValidityBitmapFailure("fewer non-null slots", /*num_values=*/3,
+                              /*null_count=*/1,
+                              /*valid_bits=*/0b00000001);
+}
+
 TEST(DictEncodingAdHoc, PutDictionaryPutIndices) {
   // Part of ARROW-3246
   auto dict_values =