diff --git a/cpp/src/parquet/encoder.cc b/cpp/src/parquet/encoder.cc index 30341bbe25e..89966b7b319 100644 --- a/cpp/src/parquet/encoder.cc +++ b/cpp/src/parquet/encoder.cc @@ -18,6 +18,7 @@ #include "parquet/encoding.h" #include +#include #include #include #include @@ -1384,9 +1385,9 @@ void DeltaBitPackEncoder::PutSpaced(const T* src, int num_values, class DeltaLengthByteArrayEncoder : public EncoderImpl, virtual public TypedEncoder { public: - explicit DeltaLengthByteArrayEncoder(const ColumnDescriptor* descr, MemoryPool* pool) - : EncoderImpl(descr, Encoding::DELTA_LENGTH_BYTE_ARRAY, - pool = ::arrow::default_memory_pool()), + explicit DeltaLengthByteArrayEncoder(const ColumnDescriptor* descr, + MemoryPool* pool = ::arrow::default_memory_pool()) + : EncoderImpl(descr, Encoding::DELTA_LENGTH_BYTE_ARRAY, pool), sink_(pool), length_encoder_(nullptr, pool) {} @@ -1408,6 +1409,67 @@ class DeltaLengthByteArrayEncoder : public EncoderImpl, protected: template void PutBinaryArray(const ArrayType& array) { + if (array.length() == 0 || array.null_count() == array.length()) { + return; + } + + constexpr int kBatchSize = 256; + std::array lengths; + + if constexpr (!std::is_same_v) { + if (array.null_count() == 0) { + const int64_t num_values = array.length(); + const auto* offsets = array.raw_value_offsets(); + const int64_t start_offset = offsets[0]; + const int64_t total_bytes = + static_cast(offsets[num_values]) - start_offset; + if constexpr (sizeof(typename ArrayType::offset_type) > sizeof(int32_t)) { + for (int64_t i = 0; i < num_values; ++i) { + const int64_t len = offsets[i + 1] - offsets[i]; + if (ARROW_PREDICT_FALSE(len < 0 || + static_cast(len) > kMaxByteArraySize)) { + PARQUET_THROW_NOT_OK(Status::Invalid( + "Parquet cannot store strings with size 2GB or more, got: ", len)); + } + } + } + if (ARROW_PREDICT_FALSE(total_bytes < 0 || + total_bytes > std::numeric_limits::max() - + sink_.length())) { + PARQUET_THROW_NOT_OK( + Status::Invalid("excess expansion in DELTA_LENGTH_BYTE_ARRAY")); + } + if (total_bytes > 0) { + PARQUET_THROW_NOT_OK(sink_.Reserve(total_bytes)); + } + for (int64_t idx = 0; idx < num_values; idx += kBatchSize) { + const int batch_size = + static_cast(std::min(kBatchSize, num_values - idx)); + for (int j = 0; j < batch_size; ++j) { + const int64_t len = + static_cast(offsets[idx + j + 1]) - offsets[idx + j]; + if (ARROW_PREDICT_FALSE(len < 0)) { + PARQUET_THROW_NOT_OK( + Status::Invalid("Negative string length in BinaryArray: ", len)); + } + lengths[j] = static_cast(len); + } + length_encoder_.Put(lengths.data(), batch_size); + } + if (total_bytes > 0) { + sink_.UnsafeAppend(array.raw_data() + start_offset, total_bytes); + unencoded_byte_array_data_bytes_ += total_bytes; + } + return; + } + } + + const int64_t estimated_bytes = GetBinaryDataSize(array); + if (estimated_bytes > 0 && + estimated_bytes <= std::numeric_limits::max() - sink_.length()) { + PARQUET_THROW_NOT_OK(sink_.Reserve(estimated_bytes)); + } + int batch_len = 0; PARQUET_THROW_NOT_OK(::arrow::VisitArraySpanInline( *array.data(), [&](::std::string_view view) { @@ -1420,12 +1482,21 @@ class DeltaLengthByteArrayEncoder : public EncoderImpl, static_cast(std::numeric_limits::max()))) { return Status::Invalid("excess expansion in DELTA_LENGTH_BYTE_ARRAY"); } - length_encoder_.Put({static_cast(view.length())}, 1); - PARQUET_THROW_NOT_OK(sink_.Append(view.data(), view.length())); - unencoded_byte_array_data_bytes_ += view.size(); + lengths[batch_len++] = static_cast(view.length()); + if (batch_len == kBatchSize) { + length_encoder_.Put(lengths.data(), kBatchSize); + batch_len = 0; + } + if (!view.empty()) { + RETURN_NOT_OK(sink_.Append(view.data(), static_cast(view.length()))); + unencoded_byte_array_data_bytes_ += view.size(); + } return Status::OK(); }, []() { return Status::OK(); })); + if (batch_len > 0) { + length_encoder_.Put(lengths.data(), batch_len); + } } ::arrow::BufferBuilder sink_; @@ -1469,11 +1540,15 @@ void DeltaLengthByteArrayEncoder::Put(const T* src, int num_values) { if (sink_.length() + total_increment_size > std::numeric_limits::max()) { throw ParquetException("excess expansion in DELTA_LENGTH_BYTE_ARRAY"); } - PARQUET_THROW_NOT_OK(sink_.Reserve(total_increment_size)); - for (int idx = 0; idx < num_values; idx++) { - sink_.UnsafeAppend(src[idx].ptr, src[idx].len); + if (total_increment_size > 0) { + PARQUET_THROW_NOT_OK(sink_.Reserve(total_increment_size)); + for (int idx = 0; idx < num_values; idx++) { + if (src[idx].len > 0) { + sink_.UnsafeAppend(src[idx].ptr, src[idx].len); + } + } + unencoded_byte_array_data_bytes_ += total_increment_size; } - unencoded_byte_array_data_bytes_ += total_increment_size; } void DeltaLengthByteArrayEncoder::PutSpaced(const T* src, int num_values, diff --git a/cpp/src/parquet/encoding_test.cc b/cpp/src/parquet/encoding_test.cc index 672e55a900d..8749314a8ca 100644 --- a/cpp/src/parquet/encoding_test.cc +++ b/cpp/src/parquet/encoding_test.cc @@ -24,6 +24,7 @@ #include #include #include +#include #include #include @@ -32,6 +33,7 @@ #include "arrow/array/builder_dict.h" #include "arrow/array/concatenate.h" #include "arrow/compute/cast.h" +#include "arrow/memory_pool.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/random.h" #include "arrow/testing/util.h" @@ -2391,6 +2393,214 @@ TEST(DeltaLengthByteArrayEncodingAdHoc, ArrowDirectPut) { } } +TEST(DeltaLengthByteArrayEncodingAdHoc, UsesProvidedMemoryPool) { + constexpr int kCount = 320; + ::arrow::StringBuilder dense_builder; + ::arrow::StringBuilder nullable_builder; + ::arrow::LargeStringBuilder large_builder; + ::arrow::StringViewBuilder view_builder; + + int64_t expected_payload_bytes = 0; + std::vector raw_strings; + raw_strings.reserve(kCount); + for (int i = 0; i < kCount; ++i) { + std::string val = (i % 2 == 0) ? ("short_" + std::to_string(i)) + : ("long_out_of_line_payload_" + std::to_string(i)); + raw_strings.push_back(val); + ASSERT_OK(dense_builder.Append(val)); + expected_payload_bytes += static_cast(val.size()); + + if (i % 7 == 0) { + ASSERT_OK(nullable_builder.AppendNull()); + } else { + ASSERT_OK(nullable_builder.Append(val)); + expected_payload_bytes += static_cast(val.size()); + } + ASSERT_OK(large_builder.Append(val)); + ASSERT_OK(view_builder.Append(val)); + expected_payload_bytes += static_cast(val.size()); + } + ASSERT_OK_AND_ASSIGN(auto dense_arr, dense_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto nullable_arr, nullable_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto large_arr, large_builder.Finish()); + auto sliced_large = large_arr->Slice(5, kCount - 10); + for (int i = 5; i < kCount - 5; ++i) { + expected_payload_bytes += static_cast(raw_strings[i].size()); + } + ASSERT_OK_AND_ASSIGN(auto view_arr, view_builder.Finish()); + + std::vector raw_ba; + raw_ba.reserve(raw_strings.size()); + for (const auto& s : raw_strings) { + raw_ba.emplace_back(static_cast(s.size()), + reinterpret_cast(s.data())); + expected_payload_bytes += static_cast(s.size()); + } + + ::arrow::ProxyMemoryPool tracking_pool(default_memory_pool()); + { + auto encoder = MakeTypedEncoder(Encoding::DELTA_LENGTH_BYTE_ARRAY, + /*use_dictionary=*/false, + /*descr=*/nullptr, &tracking_pool); + EXPECT_EQ(&tracking_pool, encoder->memory_pool()); + + for (const auto& arr : {dense_arr, nullable_arr, sliced_large, view_arr}) { + encoder->Put(*arr); + auto flushed = encoder->FlushValues(); + ASSERT_GT(flushed->size(), 0); + } + encoder->Put(raw_ba.data(), static_cast(raw_ba.size())); + auto flushed = encoder->FlushValues(); + ASSERT_GT(flushed->size(), 0); + + auto dba_encoder = MakeTypedEncoder(Encoding::DELTA_BYTE_ARRAY, + /*use_dictionary=*/false, + /*descr=*/nullptr, &tracking_pool); + dba_encoder->Put(*dense_arr); + ASSERT_GT(dba_encoder->FlushValues()->size(), 0); + } + + EXPECT_GT(tracking_pool.num_allocations(), 0); + EXPECT_GE(tracking_pool.total_bytes_allocated(), expected_payload_bytes); + EXPECT_EQ(0, tracking_pool.bytes_allocated()); +} + +TEST(DeltaLengthByteArrayEncodingAdHoc, BatchedLengthsSlicesAndZeroLengthSafety) { + auto descr = ExampleDescr(); + auto encoder = MakeTypedEncoder(Encoding::DELTA_LENGTH_BYTE_ARRAY, + /*use_dictionary=*/false, descr.get()); + auto decoder = + MakeTypedDecoder(Encoding::DELTA_LENGTH_BYTE_ARRAY, descr.get()); + + auto VerifyArrayRoundTrip = [&](const std::shared_ptr<::arrow::Array>& arr, + const std::vector& expected_valid_lengths, + const std::string& expected_payload) { + encoder->Put(*arr); + EXPECT_EQ(static_cast(expected_payload.size()), + encoder->ReportUnencodedDataBytes()); + auto buf = encoder->FlushValues(); + + auto expected_len_buf = DeltaEncode(expected_valid_lengths); + ASSERT_EQ(expected_len_buf->size() + static_cast(expected_payload.size()), + buf->size()); + EXPECT_TRUE(SliceBuffer(buf, 0, expected_len_buf->size())->Equals(*expected_len_buf)); + EXPECT_TRUE( + SliceBuffer(buf, expected_len_buf->size())->Equals(Buffer(expected_payload))); + + const int num_valid = static_cast(expected_valid_lengths.size()); + decoder->SetData(num_valid, buf->data(), static_cast(buf->size())); + typename EncodingTraits::Accumulator acc; + ASSERT_OK_AND_ASSIGN(acc.builder, ::arrow::MakeBuilder(arr->type())); + ASSERT_EQ(num_valid, + decoder->DecodeArrow(static_cast(arr->length()), + static_cast(arr->null_count()), + arr->null_bitmap_data(), arr->offset(), &acc)); + std::shared_ptr<::arrow::Array> decoded; + ASSERT_OK(acc.builder->Finish(&decoded)); + ::arrow::AssertArraysEqual(*arr, *decoded, /*verbose=*/true); + }; + + for (int count : {0, 1, 255, 256, 257, 512, 600}) { + ARROW_SCOPED_TRACE("count = ", count); + ::arrow::StringBuilder dense_b; + ::arrow::StringBuilder nullable_b; + ::arrow::StringBuilder empty_b; + ::arrow::StringBuilder nulls_b; + std::vector dense_lens, nullable_lens, empty_lens(count, 0); + std::string dense_payload, nullable_payload; + + for (int i = 0; i < count; ++i) { + std::string s = (i % 11 == 0) ? "" : ("v_" + std::to_string(i)); + ASSERT_OK(dense_b.Append(s)); + dense_lens.push_back(static_cast(s.size())); + dense_payload.append(s); + + if (i % 5 == 0) { + ASSERT_OK(nullable_b.AppendNull()); + } else { + ASSERT_OK(nullable_b.Append(s)); + nullable_lens.push_back(static_cast(s.size())); + nullable_payload.append(s); + } + ASSERT_OK(empty_b.Append("")); + ASSERT_OK(nulls_b.AppendNull()); + } + ASSERT_OK_AND_ASSIGN(auto dense_arr, dense_b.Finish()); + ASSERT_OK_AND_ASSIGN(auto nullable_arr, nullable_b.Finish()); + ASSERT_OK_AND_ASSIGN(auto empty_arr, empty_b.Finish()); + ASSERT_OK_AND_ASSIGN(auto nulls_arr, nulls_b.Finish()); + + VerifyArrayRoundTrip(dense_arr, dense_lens, dense_payload); + VerifyArrayRoundTrip(nullable_arr, nullable_lens, nullable_payload); + VerifyArrayRoundTrip(empty_arr, empty_lens, ""); + VerifyArrayRoundTrip(nulls_arr, {}, ""); + } + + { + ::arrow::LargeStringBuilder large_b; + std::vector vals; + for (int i = 0; i < 300; ++i) { + vals.push_back((i % 13 == 0) ? "" : ("slice_item_" + std::to_string(i))); + ASSERT_OK(large_b.Append(vals.back())); + } + ASSERT_OK_AND_ASSIGN(auto large_arr, large_b.Finish()); + auto sliced = large_arr->Slice(17, 260); + std::vector lens; + std::string payload; + for (int i = 17; i < 17 + 260; ++i) { + lens.push_back(static_cast(vals[i].size())); + payload.append(vals[i]); + } + VerifyArrayRoundTrip(sliced, lens, payload); + } + + { + ::arrow::StringViewBuilder view_b; + std::vector lens; + std::string payload; + for (int i = 0; i < 270; ++i) { + if (i % 9 == 0) { + ASSERT_OK(view_b.AppendNull()); + continue; + } + std::string s = (i % 3 == 0) ? "" + : (i % 3 == 1) ? ("in_" + std::to_string(i)) + : ("out_of_line_string_view_" + std::to_string(i)); + ASSERT_OK(view_b.Append(s)); + lens.push_back(static_cast(s.size())); + payload.append(s); + } + ASSERT_OK_AND_ASSIGN(auto view_arr, view_b.Finish()); + VerifyArrayRoundTrip(view_arr, lens, payload); + } + + { + const std::string abc = "abc"; + std::vector mixed = { + ByteArray(0, nullptr), + ByteArray(3, reinterpret_cast(abc.data())), + ByteArray(0, nullptr), + }; + encoder->Put(mixed.data(), static_cast(mixed.size())); + EXPECT_EQ(3, encoder->ReportUnencodedDataBytes()); + auto buf = encoder->FlushValues(); + EXPECT_EQ(DeltaEncode({0, 3, 0})->ToString() + "abc", buf->ToString()); + + std::vector all_zero_null(4, ByteArray(0, nullptr)); + encoder->Put(all_zero_null.data(), static_cast(all_zero_null.size())); + EXPECT_EQ(0, encoder->ReportUnencodedDataBytes()); + auto zero_buf = encoder->FlushValues(); + EXPECT_EQ(DeltaEncode({0, 0, 0, 0})->ToString(), zero_buf->ToString()); + + const uint8_t valid_bits = 0b00001010; + encoder->PutSpaced(all_zero_null.data(), 4, &valid_bits, /*valid_bits_offset=*/0); + EXPECT_EQ(0, encoder->ReportUnencodedDataBytes()); + auto spaced_buf = encoder->FlushValues(); + EXPECT_EQ(DeltaEncode(std::vector{0, 0})->ToString(), + spaced_buf->ToString()); + } +} + // ---------------------------------------------------------------------- // DELTA_BYTE_ARRAY encode/decode tests.