diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 4bd08fd6530..08b3a8eaddf 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -710,6 +710,9 @@ class RleBitPackedEncoder { /// This value must be representable with bit_width_ bits. bool Put(uint64_t value); + /// Encode the same value up to count times and return how many were accepted + int PutRepeated(uint64_t value, int count); + /// Flushes any pending values to the underlying buffer. /// Returns the total number of bytes written int Flush(); @@ -1526,6 +1529,28 @@ inline bool RleBitPackedEncoder::Put(uint64_t value) { return true; } +inline int RleBitPackedEncoder::PutRepeated(uint64_t value, int count) { + ARROW_DCHECK_GE(count, 0); + ARROW_DCHECK(bit_width_ == 64 || value < (uint64_t{1} << bit_width_)); + + int encoded = 0; + while (encoded < count && !buffer_full_) { + if (current_value_ == value && repeat_count_ > 8) { + // Extend an existing repeat by adding to its count instead of visiting each value + const int added = + std::min(count - encoded, std::numeric_limits::max() - repeat_count_); + repeat_count_ += added; + encoded += added; + break; + } + if (!Put(value)) { + break; + } + ++encoded; + } + return encoded; +} + inline void RleBitPackedEncoder::FlushLiteralRun(bool update_indicator_byte) { if (literal_indicator_byte_ == NULL) { // The literal indicator byte has not been reserved yet, get one now. diff --git a/cpp/src/arrow/util/rle_encoding_test.cc b/cpp/src/arrow/util/rle_encoding_test.cc index 16b702c23d4..16f7bd9b9b6 100644 --- a/cpp/src/arrow/util/rle_encoding_test.cc +++ b/cpp/src/arrow/util/rle_encoding_test.cc @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -1453,4 +1454,74 @@ TEST(RleBitPacked, GetBatchSpacedRoundtripUint64) { DoTestGetBatchSpacedRoundtrip(); } +// Check that writing repeated values together gives the same bytes as one at a time +TEST(RleBitPacked, PutRepeatedMatchesPut) { + for (int width : {0, 1, 2, 3, 8, 16, 32, 64}) { + const uint64_t maximum = + width == 64 ? std::numeric_limits::max() : ((uint64_t{1} << width) - 1); + for (uint64_t value : {uint64_t{0}, maximum}) { + for (int prefix : {0, 1, 7, 8, 9, 503, 504, 505}) { + for (int count : {0, 1, 7, 8, 9, 63, 64, 767, 768, 4096}) { + const int total = prefix + count + 9; + const int capacity = + static_cast(RleBitPackedEncoder::MaxBufferSize(width, total) + + RleBitPackedEncoder::MinBufferSize(width)); + std::vector expected(capacity); + std::vector actual(capacity); + RleBitPackedEncoder reference(expected.data(), capacity, width); + RleBitPackedEncoder encoder(actual.data(), capacity, width); + for (int i = 0; i < prefix; ++i) { + const uint64_t item = static_cast(i & 1) & maximum; + ASSERT_TRUE(reference.Put(item)); + ASSERT_TRUE(encoder.Put(item)); + } + for (int i = 0; i < count; ++i) { + ASSERT_TRUE(reference.Put(value)); + } + const int first = count / 2; + ASSERT_EQ(first, encoder.PutRepeated(value, first)); + ASSERT_EQ(count - first, encoder.PutRepeated(value, count - first)); + for (int i = 0; i < 9; ++i) { + const uint64_t item = static_cast(i & 1) & maximum; + ASSERT_TRUE(reference.Put(item)); + ASSERT_TRUE(encoder.Put(item)); + } + expected.resize(reference.Flush()); + actual.resize(encoder.Flush()); + ASSERT_EQ(expected, actual); + } + } + } + } +} + +// Check that both ways of writing stop at the same point when the buffer fills up +TEST(RleBitPacked, PutRepeatedFullBufferMatchesPut) { + for (int width : {1, 3, 16}) { + const int capacity = static_cast(RleBitPackedEncoder::MinBufferSize(width)); + std::vector expected(capacity); + std::vector actual(capacity); + RleBitPackedEncoder reference(expected.data(), capacity, width); + RleBitPackedEncoder encoder(actual.data(), capacity, width); + bool stopped = false; + for (int run = 0; run < 4096; ++run) { + const uint64_t value = run & 1; + const int count = run % 17 + 1; + int accepted = 0; + while (accepted < count && reference.Put(value)) { + ++accepted; + } + ASSERT_EQ(accepted, encoder.PutRepeated(value, count)); + if (accepted != count) { + stopped = true; + break; + } + } + ASSERT_TRUE(stopped); + expected.resize(reference.Flush()); + actual.resize(encoder.Flush()); + ASSERT_EQ(expected, actual); + } +} + } // namespace arrow::util diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index d858962d472..a60a8217987 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -42,6 +42,7 @@ #include "arrow/compute/api.h" #include "arrow/extension/json.h" #include "arrow/io/api.h" +#include "arrow/memory_pool.h" #include "arrow/record_batch.h" #include "arrow/scalar.h" #include "arrow/table.h" @@ -3531,6 +3532,234 @@ TEST(ArrowReadWrite, EmptyListView) { ASSERT_EQ(0, list_view.value_sizes()->size()); } +// Build the same numbers as fixed-size and ordinary lists for comparison +template +static Result, std::shared_ptr>> +MakeFixedListTestTables(int32_t list_size, bool root_nullable, bool item_nullable, + int64_t rows = 65, bool chunked = true) { + ::arrow::NumericBuilder values_builder; + const int64_t value_count = (rows + 2) * list_size + 3; + RETURN_NOT_OK(values_builder.Reserve(value_count)); + for (int64_t i = 0; i < value_count; ++i) { + RETURN_NOT_OK(values_builder.Append( + static_cast((i * 17) % 1009 - 504))); + } + ARROW_ASSIGN_OR_RAISE(auto values, values_builder.Finish()); + values = values->Slice(3, (rows + 2) * list_size); + ::arrow::Int32Builder offsets_builder; + for (int64_t i = 0; i <= rows + 2; ++i) { + RETURN_NOT_OK(offsets_builder.Append(static_cast(i * list_size))); + } + ARROW_ASSIGN_OR_RAISE(auto offsets, offsets_builder.Finish()); + auto item = ::arrow::field("item", std::make_shared(), item_nullable); + ARROW_ASSIGN_OR_RAISE(auto fixed, + ::arrow::FixedSizeListArray::FromArrays( + values, ::arrow::fixed_size_list(item, list_size))); + ARROW_ASSIGN_OR_RAISE(auto ordinary, ::arrow::ListArray::FromArrays(::arrow::list(item), + *offsets, *values)); + auto make_table = [&](const std::shared_ptr& array) { + auto sliced = array->Slice(1, rows); + ArrayVector chunks; + if (chunked) { + const int64_t first = rows / 3; + const int64_t second = rows / 3; + chunks = {sliced->Slice(0, 0), sliced->Slice(0, first), sliced->Slice(first, 0), + sliced->Slice(first, second), + sliced->Slice(first + second, rows - first - second)}; + } else { + chunks = {sliced}; + } + auto column = std::make_shared(chunks); + auto schema = + ::arrow::schema({::arrow::field("root", sliced->type(), root_nullable)}); + return Table::Make(schema, {column}); + }; + return std::make_pair(make_table(fixed), make_table(ordinary)); +} + +// Write a test table into memory so we can compare its saved bytes +static Result> WriteFixedListTestTable( + const Table& table, const std::shared_ptr& properties, + int64_t row_group_size, bool store_schema = false, bool buffered = false) { + ArrowWriterProperties::Builder arrow_properties; + arrow_properties.set_use_threads(false); + if (store_schema) { + arrow_properties.store_schema(); + } + ARROW_ASSIGN_OR_RAISE(auto sink, ::arrow::io::BufferOutputStream::Create()); + if (buffered) { + ARROW_ASSIGN_OR_RAISE(auto writer, + FileWriter::Open(*table.schema(), default_memory_pool(), sink, + properties, arrow_properties.build())); + RETURN_NOT_OK(writer->NewBufferedRowGroup()); + ::arrow::TableBatchReader batches(table); + batches.set_chunksize(7); + while (true) { + std::shared_ptr<::arrow::RecordBatch> batch; + RETURN_NOT_OK(batches.ReadNext(&batch)); + if (!batch) { + break; + } + RETURN_NOT_OK(writer->WriteRecordBatch(*batch)); + } + RETURN_NOT_OK(writer->Close()); + } else { + RETURN_NOT_OK(WriteTable(table, default_memory_pool(), sink, row_group_size, + properties, arrow_properties.build())); + } + return sink->Finish(); +} + +// Compare saved bytes and read-back data across different writer settings +template +static void CheckFixedSizeListFastPath() { + struct Settings { + int64_t batch_size; + int64_t page_size; + int64_t rows_per_page; + int64_t row_group_size; + int64_t dictionary_size; + }; + const std::vector settings = { + {1, 64, 1, 1, -1}, + {7, 128, 3, 17, 64}, + {1024, 1 << 20, 4096, 65, 1 << 20}, + }; + for (int32_t list_size : {1, 3, 8, 9, 16, 768}) { + for (bool root_nullable : {false, true}) { + for (bool item_nullable : {false, true}) { + ASSERT_OK_AND_ASSIGN(auto tables, MakeFixedListTestTables( + list_size, root_nullable, item_nullable)); + const auto& fixed_table = tables.first; + const auto& list_table = tables.second; + for (auto version : {ParquetDataPageVersion::V1, ParquetDataPageVersion::V2}) { + for (bool page_index : {false, true}) { + for (const auto& setting : settings) { + SCOPED_TRACE(::testing::Message() + << "list_size=" << list_size << " root_nullable=" + << root_nullable << " item_nullable=" << item_nullable + << " version=" << static_cast(version) << " page_index=" + << page_index << " batch_size=" << setting.batch_size); + WriterProperties::Builder builder; + builder.data_page_version(version); + builder.write_batch_size(setting.batch_size); + builder.data_pagesize(setting.page_size); + builder.max_rows_per_page(setting.rows_per_page); + builder.set_size_statistics_level(SizeStatisticsLevel::PageAndColumnChunk); + if (page_index) { + builder.enable_write_page_index(); + } else { + builder.disable_write_page_index(); + } + if (setting.dictionary_size < 0) { + builder.disable_dictionary(); + } else { + builder.enable_dictionary(); + builder.dictionary_pagesize_limit(setting.dictionary_size); + } + auto properties = builder.build(); + ASSERT_OK_AND_ASSIGN(auto expected, + WriteFixedListTestTable(*list_table, properties, + setting.row_group_size)); + ASSERT_OK_AND_ASSIGN(auto actual, + WriteFixedListTestTable(*fixed_table, properties, + setting.row_group_size)); + ASSERT_TRUE(expected->Equals(*actual)); + ASSERT_OK_AND_ASSIGN(auto stored, + WriteFixedListTestTable(*fixed_table, properties, + setting.row_group_size, true)); + ASSERT_OK_AND_ASSIGN(auto reader, + OpenFile(std::make_shared(stored), + default_memory_pool())); + ASSERT_OK_AND_ASSIGN(auto restored, reader->ReadTable()); + ASSERT_TRUE(fixed_table->Equals(*restored)); + } + } + } + } + } + } +} + +// Check that the shortcut saves identical files and reads back the original lists +TEST(ArrowReadWrite, FixedSizeListFastPathMatchesList) { + ASSERT_NO_FATAL_FAILURE(CheckFixedSizeListFastPath<::arrow::Int32Type>()); + ASSERT_NO_FATAL_FAILURE(CheckFixedSizeListFastPath<::arrow::Int64Type>()); + ASSERT_NO_FATAL_FAILURE(CheckFixedSizeListFastPath<::arrow::FloatType>()); + ASSERT_NO_FATAL_FAILURE(CheckFixedSizeListFastPath<::arrow::DoubleType>()); +} + +// Measure the highest amount of writer memory tracked by Arrow during a write +static Result MeasureFixedListWriterMemory(const Table& table) { + ::arrow::ProxyMemoryPool pool(default_memory_pool()); + { + ARROW_ASSIGN_OR_RAISE(auto sink, ::arrow::io::BufferOutputStream::Create()); + WriterProperties::Builder properties; + properties.memory_pool(&pool); + properties.disable_dictionary(); + properties.compression(Compression::UNCOMPRESSED); + properties.write_batch_size(table.num_rows() * 3); + properties.data_pagesize(64 << 20); + properties.max_rows_per_page(table.num_rows() + 1); + properties.set_size_statistics_level(SizeStatisticsLevel::None); + ArrowWriterProperties::Builder arrow_properties; + arrow_properties.set_use_threads(false); + RETURN_NOT_OK(WriteTable(table, &pool, sink, table.num_rows(), properties.build(), + arrow_properties.build())); + } + return pool.max_memory(); +} + +// Check that the shortcut saves at least one marker array worth of temporary memory +TEST(ArrowReadWrite, FixedSizeListFastPathUsesLessMemory) { + constexpr int64_t rows = 65536; + constexpr int32_t list_size = 3; + ASSERT_OK_AND_ASSIGN(auto tables, MakeFixedListTestTables<::arrow::FloatType>( + list_size, false, false, rows, false)); + ASSERT_OK_AND_ASSIGN(const int64_t fast_peak, + MeasureFixedListWriterMemory(*tables.first)); + ASSERT_OK_AND_ASSIGN(const int64_t generic_peak, + MeasureFixedListWriterMemory(*tables.second)); + const int64_t one_marker_array = + rows * list_size * static_cast(sizeof(int16_t)); + ASSERT_LE(fast_peak + one_marker_array, generic_peak); +} + +// Check that missing lists, missing values and unsupported settings keep correct output +TEST(ArrowReadWrite, FixedSizeListFallbackMatchesList) { + const std::vector> cases = { + {R"([{"root": [1, 2, 3]}])", R"([{"root": [4, 5, 6]}, {"root": [7, 8, 9]}])"}, + {R"([{"root": [1, 2, 3]}])", R"([{"root": null}, {"root": [4, null, 6]}])"}, + }; + auto item = ::arrow::field("item", ::arrow::int32()); + auto fixed_schema = + ::arrow::schema({::arrow::field("root", ::arrow::fixed_size_list(item, 3))}); + auto list_schema = ::arrow::schema({::arrow::field("root", ::arrow::list(item))}); + for (const auto& json : cases) { + auto fixed_table = ::arrow::TableFromJSON(fixed_schema, json); + auto list_table = ::arrow::TableFromJSON(list_schema, json); + for (auto version : {ParquetDataPageVersion::V1, ParquetDataPageVersion::V2}) { + for (bool buffered : {false, true}) { + for (bool content_defined : {false, true}) { + WriterProperties::Builder builder; + builder.data_page_version(version); + if (content_defined) { + builder.enable_content_defined_chunking(); + } + auto properties = builder.build(); + ASSERT_OK_AND_ASSIGN( + auto expected, + WriteFixedListTestTable(*list_table, properties, 65, false, buffered)); + ASSERT_OK_AND_ASSIGN( + auto actual, + WriteFixedListTestTable(*fixed_table, properties, 65, false, buffered)); + ASSERT_TRUE(expected->Equals(*actual)); + } + } + } + } +} + struct FixedSizeListTestCase { std::shared_ptr type; std::string json; diff --git a/cpp/src/parquet/arrow/reader_writer_benchmark.cc b/cpp/src/parquet/arrow/reader_writer_benchmark.cc index 2f288fd2eb0..0471f90ea40 100644 --- a/cpp/src/parquet/arrow/reader_writer_benchmark.cc +++ b/cpp/src/parquet/arrow/reader_writer_benchmark.cc @@ -222,6 +222,61 @@ BENCHMARK_TEMPLATE2(BM_WriteColumn, true, DoubleType); BENCHMARK_TEMPLATE2(BM_WriteColumn, false, BooleanType); BENCHMARK_TEMPLATE2(BM_WriteColumn, true, BooleanType); +// Measure writing the same numbers as a flat column or fixed-size lists +static void BM_WriteFloatLayout(::benchmark::State& state) { + const bool use_lists = state.range(0) == 1; + const int32_t list_size = static_cast(state.range(1)); + const auto page_version = + state.range(2) == 1 ? ParquetDataPageVersion::V1 : ParquetDataPageVersion::V2; + const int64_t num_values = (BENCHMARK_SIZE / list_size) * list_size; + + ::arrow::random::RandomArrayGenerator rng(500); + auto values = rng.Float32(num_values, -1.0f, 1.0f, 0.0); + std::shared_ptr<::arrow::Array> column = values; + + if (use_lists) { + auto item_field = ::arrow::field("item", ::arrow::float32(), false); + auto list_type = ::arrow::fixed_size_list(item_field, list_size); + auto lists = ::arrow::FixedSizeListArray::FromArrays(values, list_type); + EXIT_NOT_OK(lists.status()); + column = *lists; + } + + auto column_field = ::arrow::field("column", column->type(), false); + auto schema = ::arrow::schema({column_field}); + auto table = ::arrow::Table::Make(schema, {column}); + + WriterProperties::Builder builder; + builder.disable_dictionary(); + builder.compression(Compression::UNCOMPRESSED); + builder.data_page_version(page_version); + auto properties = builder.build(); + + state.SetLabel(use_lists ? "fixed_size_list" : "flat"); + while (state.KeepRunning()) { + auto output = CreateOutputStream(); + EXIT_NOT_OK(WriteTable(*table, ::arrow::default_memory_pool(), output, + table->num_rows(), properties)); + } + + const int64_t total_values = state.iterations() * num_values; + state.SetItemsProcessed(total_values); + state.SetBytesProcessed(total_values * sizeof(float)); +} + +// Run both layouts with six list sizes and both Parquet page versions +static void FloatLayoutWriteArguments(::benchmark::internal::Benchmark* b) { + for (int64_t layout : {0, 1}) { + for (int64_t list_size : {1, 3, 8, 9, 16, 768}) { + for (int64_t page_version : {1, 2}) { + b->Args({layout, list_size, page_version}); + } + } + } +} + +BENCHMARK(BM_WriteFloatLayout)->Apply(FloatLayoutWriteArguments); + int32_t kInfiniteUniqueValues = -1; std::shared_ptr
RandomStringTable(std::shared_ptr<::arrow::DataType> type, diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index e0fbe308219..90380f30ef8 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -135,6 +135,42 @@ Status GetSchemaMetadata(const ::arrow::Schema& schema, ::arrow::MemoryPool* poo return Status::OK(); } +// Check that all chunks contain supported fixed-size lists with no missing values +bool CanWriteFixedSizeLists(const std::vector>& arrays) { + if (arrays.empty() || arrays.front()->type_id() != ::arrow::Type::FIXED_SIZE_LIST) { + return false; + } + const auto& type = + checked_cast(*arrays.front()->type()); + if (type.list_size() <= 0) { + return false; + } + switch (type.value_type()->id()) { + case ::arrow::Type::INT32: + case ::arrow::Type::INT64: + case ::arrow::Type::FLOAT: + case ::arrow::Type::DOUBLE: + break; + default: + return false; + } + for (const auto& array : arrays) { + if (!array->type()->Equals(arrays.front()->type())) { + return false; + } + const auto& lists = checked_cast(*array); + if (lists.null_count() != 0) { + return false; + } + auto values = + lists.values()->Slice(lists.value_offset(0), lists.length() * type.list_size()); + if (values->null_count() != 0) { + return false; + } + } + return true; +} + // Manages writing nested parquet columns with support for all nested types // supported by parquet. class ArrowColumnWriterV2 { @@ -165,6 +201,30 @@ class ArrowColumnWriterV2 { } else { PARQUET_CATCH_NOT_OK(column_writer = row_group_writer_->NextColumn()); } + if (!fixed_size_lists_.empty()) { + if (column_writer->properties()->content_defined_chunking_enabled()) { + for (const auto& array : fixed_size_lists_) { + ARROW_ASSIGN_OR_RAISE(auto builder, + MultipathLevelBuilder::Make(*array, is_nullable_)); + level_builders_.emplace_back(std::move(builder)); + } + fixed_size_lists_.clear(); + } else { + const auto& type = checked_cast( + *fixed_size_lists_.front()->type()); + auto fast_ctx = *ctx; + fast_ctx.fixed_size_list_size = type.list_size(); + for (const auto& array : fixed_size_lists_) { + const auto& lists = checked_cast(*array); + // Take only the numbers belonging to this slice of lists + auto values = lists.values()->Slice(lists.value_offset(0), + lists.length() * type.list_size()); + RETURN_NOT_OK(column_writer->WriteArrow(nullptr, nullptr, values->length(), + *values, &fast_ctx, + type.value_field()->nullable())); + } + } + } for (auto& level_builder : level_builders_) { RETURN_NOT_OK(level_builder->Write( leaf_idx, ctx, [&](const MultipathLevelBuilderResult& result) { @@ -241,6 +301,7 @@ class ArrowColumnWriterV2 { int64_t values_written = 0; std::vector> builders; + std::vector> arrays; const int leaf_count = CalculateLeafCount(data.type().get()); bool is_nullable = false; @@ -283,13 +344,7 @@ class ArrowColumnWriterV2 { std::shared_ptr array_to_write = chunk.Slice(chunk_offset, chunk_write_size); if (array_to_write->length() > 0) { - ARROW_ASSIGN_OR_RAISE(std::unique_ptr builder, - MultipathLevelBuilder::Make(*array_to_write, is_nullable)); - if (leaf_count != builder->GetLeafCount()) { - return Status::UnknownError("data type leaf_count != builder_leaf_count", - leaf_count, " ", builder->GetLeafCount()); - } - builders.emplace_back(std::move(builder)); + arrays.emplace_back(std::move(array_to_write)); } if (chunk_write_size == available_values) { @@ -298,8 +353,24 @@ class ArrowColumnWriterV2 { } values_written += chunk_write_size; } - return std::make_unique(std::move(builders), column_index, - leaf_count, row_group_writer); + // Check every chunk before choosing one writing path for the column + if (row_group_writer->buffered() || !CanWriteFixedSizeLists(arrays)) { + for (const auto& array : arrays) { + ARROW_ASSIGN_OR_RAISE(auto builder, + MultipathLevelBuilder::Make(*array, is_nullable)); + if (leaf_count != builder->GetLeafCount()) { + return Status::UnknownError("data type leaf_count != builder_leaf_count", + leaf_count, " ", builder->GetLeafCount()); + } + builders.emplace_back(std::move(builder)); + } + arrays.clear(); + } + auto writer = std::make_unique(std::move(builders), column_index, + leaf_count, row_group_writer); + writer->fixed_size_lists_ = std::move(arrays); + writer->is_nullable_ = is_nullable; + return writer; } int leaf_count() const { return leaf_count_; } @@ -307,6 +378,8 @@ class ArrowColumnWriterV2 { private: // One builder per column-chunk. std::vector> level_builders_; + std::vector> fixed_size_lists_; + bool is_nullable_ = false; int start_leaf_column_index_; int leaf_count_; RowGroupWriter* row_group_writer_; diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index 3b4a873b519..89b9bfd892f 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -242,6 +242,43 @@ int LevelEncoder::Encode(int batch_size, const int16_t* levels) { return num_encoded; } +// Encode a repeated marker such as "present" without first building a marker array +int LevelEncoder::EncodeRepeated(int batch_size, int16_t level) { + if (!rle_encoder_ || encoding_ != Encoding::RLE || batch_size < 0 || level < 0 || + static_cast(level) >= (uint64_t{1} << bit_width_)) { + throw ParquetException("Invalid input for repeated level encoding"); + } + const int encoded = rle_encoder_->PutRepeated(level, batch_size); + rle_length_ = rle_encoder_->Flush(); + return encoded; +} + +// Encode 0 for each list start and 1 for the remaining values in that list +int LevelEncoder::EncodeFixedSizeList(int batch_size, int32_t list_size) { + if (!rle_encoder_ || encoding_ != Encoding::RLE || bit_width_ != 1 || batch_size < 0 || + list_size <= 0) { + throw ParquetException("Invalid input for fixed-size-list level encoding"); + } + if (list_size == 1) { + return EncodeRepeated(batch_size, 0); + } + int encoded = 0; + while (encoded < batch_size) { + if (!rle_encoder_->Put(0)) { + break; + } + ++encoded; + const int count = std::min(list_size - 1, batch_size - encoded); + const int accepted = rle_encoder_->PutRepeated(1, count); + encoded += accepted; + if (accepted != count) { + break; + } + } + rle_length_ = rle_encoder_->Flush(); + return encoded; +} + // ---------------------------------------------------------------------- // PageWriter implementation @@ -837,9 +874,10 @@ class ColumnWriterImpl { repetition_levels_sink_.Append(levels, sizeof(int16_t) * num_levels)); } - // RLE encode the src_buffer into dest_buffer and return the encoded size + // Encode marker arrays or known fixed-size-list patterns into page bytes int64_t RleEncodeLevels(const void* src_buffer, ResizableBuffer* dest_buffer, - int16_t max_level, bool include_length_prefix = true); + int16_t max_level, bool include_length_prefix = true, + bool is_repetition = false); // Serialize the buffered Data Pages void FlushBufferedDataPages(); @@ -861,6 +899,8 @@ class ColumnWriterImpl { LevelEncoder level_encoder_; + int32_t fixed_size_list_size_ = 0; + MemoryPool* allocator_; // The total number of values stored in the data page. This is the maximum of @@ -935,25 +975,35 @@ class ColumnWriterImpl { // return the size of the encoded buffer int64_t ColumnWriterImpl::RleEncodeLevels(const void* src_buffer, ResizableBuffer* dest_buffer, int16_t max_level, - bool include_length_prefix) { - // V1 DataPage includes the length of the RLE level as a prefix. - int32_t prefix_size = include_length_prefix ? sizeof(int32_t) : 0; - - // TODO: This only works with due to some RLE specifics - int64_t rle_size = LevelEncoder::MaxBufferSize(Encoding::RLE, max_level, - static_cast(num_buffered_values_)) + - prefix_size; - - // Use Arrow::Buffer::shrink_to_fit = false - // underlying buffer only keeps growing. Resize to a smaller size does not reallocate. - PARQUET_THROW_NOT_OK(dest_buffer->Resize(rle_size, false)); - - level_encoder_.Init(Encoding::RLE, max_level, static_cast(num_buffered_values_), + bool include_length_prefix, + bool is_repetition) { + if (num_buffered_values_ < 0 || + num_buffered_values_ > std::numeric_limits::max()) { + throw ParquetException("Too many levels in a data page"); + } + const int count = static_cast(num_buffered_values_); + const int32_t prefix_size = + include_length_prefix ? static_cast(sizeof(int32_t)) : 0; + const int64_t capacity = + static_cast(LevelEncoder::MaxBufferSize(Encoding::RLE, max_level, count)) + + prefix_size; + PARQUET_THROW_NOT_OK(dest_buffer->Resize(capacity, false)); + level_encoder_.Init(Encoding::RLE, max_level, count, dest_buffer->mutable_data() + prefix_size, static_cast(dest_buffer->size() - prefix_size)); - int encoded = level_encoder_.Encode(static_cast(num_buffered_values_), - reinterpret_cast(src_buffer)); - DCHECK_EQ(encoded, num_buffered_values_); + int encoded; + if (fixed_size_list_size_ > 0) { + if (is_repetition) { + encoded = level_encoder_.EncodeFixedSizeList(count, fixed_size_list_size_); + } else { + encoded = level_encoder_.EncodeRepeated(count, max_level); + } + } else { + encoded = level_encoder_.Encode(count, reinterpret_cast(src_buffer)); + } + if (encoded != count) { + throw ParquetException("Could not encode all levels in a data page"); + } if (include_length_prefix) { reinterpret_cast(dest_buffer->mutable_data())[0] = level_encoder_.len(); @@ -978,7 +1028,8 @@ void ColumnWriterImpl::AddDataPage() { if (descr_->max_repetition_level() > 0) { repetition_levels_rle_size = RleEncodeLevels( repetition_levels_sink_.data(), repetition_levels_rle_.get(), - descr_->max_repetition_level(), /*include_length_prefix=*/is_v1_data_page); + descr_->max_repetition_level(), /*include_length_prefix=*/is_v1_data_page, + /*is_repetition=*/true); } int64_t uncompressed_size = @@ -1235,6 +1286,41 @@ inline void DoInBatchesRepeated(const int16_t* def_levels, const int16_t* rep_le } } +// Use the list size to find batch boundaries while keeping the existing page breaks +template +inline void DoInBatchesFixedSizeList(int64_t num_levels, int32_t list_size, + int64_t batch_size, int64_t max_rows_per_page, + bool pages_change_on_record_boundaries, + Action&& action, + GetBufferedRows&& curr_page_buffered_rows) { + if (list_size <= 0 || num_levels < 0 || num_levels % list_size != 0 || + batch_size <= 0 || max_rows_per_page <= 0) { + throw ParquetException("Invalid fixed-size-list batch settings"); + } + int64_t offset = 0; + while (offset < num_levels) { + const int64_t remaining_rows = (num_levels - offset) / list_size; + const int64_t max_batch_size = std::min(batch_size, num_levels - offset); + const int64_t batch_rows = (max_batch_size - 1) / list_size + 1; + const int64_t available_rows = max_rows_per_page - curr_page_buffered_rows(); + if (available_rows < 0) { + throw ParquetException("Data page contains too many rows"); + } + const int64_t rows = std::min(batch_rows, available_rows); + if (rows < remaining_rows) { + const int64_t count = rows * list_size; + action(offset, count, true); + offset += count; + } else { + // Handle the last list separately to match when the existing writer ends a page + const int64_t last_row = num_levels - list_size; + action(offset, last_row - offset, true); + action(last_row, list_size, !pages_change_on_record_boundaries); + offset = num_levels; + } + } +} + template inline void DoInBatches(const int16_t* def_levels, const int16_t* rep_levels, int64_t num_levels, int64_t batch_size, int64_t max_rows_per_page, @@ -1329,6 +1415,9 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, int64_t WriteBatch(int64_t num_values, const int16_t* def_levels, const int16_t* rep_levels, const T* values) override { + if (fixed_size_list_size_ > 0) { + throw ParquetException("Cannot use WriteBatch after fixed-size-list writing"); + } if (ARROW_PREDICT_FALSE(content_defined_chunker_ != nullptr)) { throw ParquetException( "Content-defined chunking is not supported in WriteBatch() or " @@ -1363,15 +1452,25 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, // circumvent this check when writing ::arrow::DictionaryArray directly CheckDictionarySizeLimit(); }; - DoInBatches(def_levels, rep_levels, num_values, properties_->write_batch_size(), - properties_->max_rows_per_page(), pages_change_on_record_boundaries(), - WriteChunk, [this]() { return num_buffered_rows_; }); + if (fixed_size_list_size_ > 0) { + DoInBatchesFixedSizeList( + num_values, fixed_size_list_size_, properties_->write_batch_size(), + properties_->max_rows_per_page(), pages_change_on_record_boundaries(), + WriteChunk, [this]() { return num_buffered_rows_; }); + } else { + DoInBatches(def_levels, rep_levels, num_values, properties_->write_batch_size(), + properties_->max_rows_per_page(), pages_change_on_record_boundaries(), + WriteChunk, [this]() { return num_buffered_rows_; }); + } return value_offset; } void WriteBatchSpaced(int64_t num_values, const int16_t* def_levels, const int16_t* rep_levels, const uint8_t* valid_bits, int64_t valid_bits_offset, const T* values) override { + if (fixed_size_list_size_ > 0) { + throw ParquetException("Cannot use WriteBatchSpaced after fixed-size-list writing"); + } if (ARROW_PREDICT_FALSE(content_defined_chunker_ != nullptr)) { throw ParquetException( "Content-defined chunking is not supported in WriteBatch() or " @@ -1424,6 +1523,44 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, int64_t num_levels, const ::arrow::Array& leaf_array, ArrowWriteContext* ctx, bool leaf_field_nullable) override { BEGIN_PARQUET_CATCH_EXCEPTIONS + if (ctx->fixed_size_list_size < 0) { + return Status::Invalid("Fixed-size-list size must not be negative"); + } + if (ctx->fixed_size_list_size > 0) { + const int32_t list_size = ctx->fixed_size_list_size; + const auto arrow_type = leaf_array.type_id(); + const bool supported_type = (std::is_same_v && + arrow_type == ::arrow::Type::INT32) || + (std::is_same_v && + arrow_type == ::arrow::Type::INT64) || + (std::is_same_v && + arrow_type == ::arrow::Type::FLOAT) || + (std::is_same_v && + arrow_type == ::arrow::Type::DOUBLE); + if (!supported_type) { + return Status::Invalid( + "Fixed-size-list writing requires a matching numeric type"); + } + if (closed_ || content_defined_chunker_ != nullptr || + descr_->max_repetition_level() != 1 || def_levels != nullptr || + rep_levels != nullptr || num_levels < 0 || num_levels % list_size != 0 || + leaf_array.length() != num_levels || leaf_array.null_count() != 0) { + return Status::Invalid("Invalid input for fixed-size-list writing"); + } + if ((fixed_size_list_size_ == 0 && num_buffered_values_ != 0) || + (fixed_size_list_size_ != 0 && fixed_size_list_size_ != list_size)) { + return Status::Invalid("Cannot mix fixed-size-list writing paths"); + } + if (num_levels == 0) { + return Status::OK(); + } + fixed_size_list_size_ = list_size; + return WriteArrowDense(nullptr, nullptr, num_levels, leaf_array, ctx, + /*maybe_parent_nulls=*/false); + } + if (fixed_size_list_size_ != 0) { + return Status::Invalid("Cannot mix fixed-size-list writing paths"); + } // Leaf nulls are canonical when there is only a single null element after a list // and it is at the leaf. bool single_nullable_element = @@ -1669,6 +1806,21 @@ class TypedColumnWriterImpl : public ColumnWriterImpl, int64_t WriteLevels(int64_t num_levels, const int16_t* def_levels, const int16_t* rep_levels) { + if (fixed_size_list_size_ > 0) { + // Count rows and markers using the list size instead of scanning marker arrays + ARROW_DCHECK_EQ(num_levels % fixed_size_list_size_, 0); + const int64_t rows = num_levels / fixed_size_list_size_; + if (page_size_statistics_) { + auto& statistics = *page_size_statistics_; + statistics.definition_level_histogram[descr_->max_definition_level()] += + num_levels; + statistics.repetition_level_histogram[0] += rows; + statistics.repetition_level_histogram[1] += num_levels - rows; + } + rows_written_ += rows; + num_buffered_rows_ += rows; + return num_levels; + } // Update histograms now, to maximize cache efficiency. UpdateLevelHistogram(num_levels, def_levels, rep_levels); diff --git a/cpp/src/parquet/column_writer.h b/cpp/src/parquet/column_writer.h index 03817a77f31..86a3a56596a 100644 --- a/cpp/src/parquet/column_writer.h +++ b/cpp/src/parquet/column_writer.h @@ -73,6 +73,12 @@ class PARQUET_EXPORT LevelEncoder { // Encodes a batch of levels from an array and returns the number of levels encoded int Encode(int batch_size, const int16_t* levels); + // Encode one marker many times without an input array using RLE + int EncodeRepeated(int batch_size, int16_t level); + + // Encode markers from a list start such as 0, 1, 1 for lists of size 3 + int EncodeFixedSizeList(int batch_size, int32_t list_size); + int32_t len() { if (encoding_ != Encoding::RLE) { throw ParquetException("Only implemented for RLE encoding"); diff --git a/cpp/src/parquet/column_writer_test.cc b/cpp/src/parquet/column_writer_test.cc index 69ec07c1b78..4192143e36c 100644 --- a/cpp/src/parquet/column_writer_test.cc +++ b/cpp/src/parquet/column_writer_test.cc @@ -22,9 +22,12 @@ #include #include +#include "arrow/array.h" #include "arrow/io/buffered.h" #include "arrow/io/file.h" +#include "arrow/testing/builder.h" #include "arrow/testing/gtest_util.h" +#include "arrow/type.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_builders.h" #include "arrow/util/config.h" @@ -1245,6 +1248,162 @@ TEST(TestColumnWriter, RepeatedListsUpdateSpacedBug) { writer->Close(); } +// Check that encoding a marker by count gives the same bytes as encoding its array +TEST(TestLevelEncoder, EncodeRepeatedMatchesEncode) { + for (int16_t maximum : {1, 2, 3, 7}) { + for (int16_t level : {int16_t{0}, maximum}) { + for (int count : {0, 1, 7, 8, 9, 767, 768, 4096}) { + const int capacity = LevelEncoder::MaxBufferSize(Encoding::RLE, maximum, count); + std::vector levels(count, level); + std::vector expected(capacity); + std::vector actual(capacity); + LevelEncoder reference; + LevelEncoder encoder; + reference.Init(Encoding::RLE, maximum, count, expected.data(), capacity); + encoder.Init(Encoding::RLE, maximum, count, actual.data(), capacity); + ASSERT_EQ(count, reference.Encode(count, levels.data())); + ASSERT_EQ(count, encoder.EncodeRepeated(count, level)); + expected.resize(reference.len()); + actual.resize(encoder.len()); + ASSERT_EQ(expected, actual); + } + } + } +} + +// Check that invalid marker values, counts and encoder settings are rejected +TEST(TestLevelEncoder, EncodeRepeatedRejectsInvalidInput) { + LevelEncoder encoder; + EXPECT_THROW(encoder.EncodeRepeated(1, 1), ParquetException); + const int capacity = LevelEncoder::MaxBufferSize(Encoding::RLE, 1, 8); + std::vector buffer(capacity); + encoder.Init(Encoding::RLE, 1, 8, buffer.data(), capacity); + EXPECT_THROW(encoder.EncodeRepeated(-1, 1), ParquetException); + EXPECT_THROW(encoder.EncodeRepeated(1, -1), ParquetException); + EXPECT_THROW(encoder.EncodeRepeated(1, 2), ParquetException); + encoder.Init(Encoding::BIT_PACKED, 1, 8, buffer.data(), capacity); + EXPECT_THROW(encoder.EncodeRepeated(1, 1), ParquetException); +} + +// Check that list sizes produce the same marker bytes as explicit 0 and 1 arrays +TEST(TestLevelEncoder, EncodeFixedSizeListMatchesEncode) { + for (int32_t list_size : {1, 3, 8, 9, 16, 768}) { + for (int count : {0, 1, 7, 8, 9, 63, 64, 767, 768, 4097}) { + const int capacity = LevelEncoder::MaxBufferSize(Encoding::RLE, 1, count); + std::vector levels(count); + for (int i = 0; i < count; ++i) { + levels[i] = i % list_size == 0 ? 0 : 1; + } + std::vector expected(capacity); + std::vector actual(capacity); + LevelEncoder reference; + LevelEncoder encoder; + reference.Init(Encoding::RLE, 1, count, expected.data(), capacity); + encoder.Init(Encoding::RLE, 1, count, actual.data(), capacity); + ASSERT_EQ(count, reference.Encode(count, levels.data())); + ASSERT_EQ(count, encoder.EncodeFixedSizeList(count, list_size)); + expected.resize(reference.len()); + actual.resize(encoder.len()); + ASSERT_EQ(expected, actual); + } + } +} + +// Check that invalid list sizes, counts and encoder settings are rejected +TEST(TestLevelEncoder, EncodeFixedSizeListRejectsInvalidInput) { + LevelEncoder encoder; + EXPECT_THROW(encoder.EncodeFixedSizeList(3, 3), ParquetException); + const int capacity = LevelEncoder::MaxBufferSize(Encoding::RLE, 2, 8); + std::vector buffer(capacity); + encoder.Init(Encoding::RLE, 1, 8, buffer.data(), capacity); + EXPECT_THROW(encoder.EncodeFixedSizeList(-1, 3), ParquetException); + EXPECT_THROW(encoder.EncodeFixedSizeList(3, 0), ParquetException); + EXPECT_THROW(encoder.EncodeFixedSizeList(3, -1), ParquetException); + encoder.Init(Encoding::RLE, 2, 8, buffer.data(), capacity); + EXPECT_THROW(encoder.EncodeFixedSizeList(3, 3), ParquetException); +} + +class TestFixedSizeListWriter : public ::testing::Test { + protected: + void MakeWriter(Type::type physical_type) { + auto item = PrimitiveNode::Make("item", Repetition::OPTIONAL, physical_type); + auto repeated = GroupNode::Make("list", Repetition::REPEATED, {item}); + auto field = + GroupNode::Make("root", Repetition::OPTIONAL, {repeated}, ConvertedType::LIST); + schema_.Init(GroupNode::Make("schema", Repetition::REQUIRED, {field})); + properties_ = WriterProperties::Builder().disable_dictionary()->build(); + metadata_ = ColumnChunkMetaDataBuilder::Make(properties_, schema_.Column(0)); + auto pager = PageWriter::Open(CreateOutputStream(), Compression::UNCOMPRESSED, + metadata_.get()); + writer_ = ColumnWriter::Make(metadata_.get(), std::move(pager), properties_.get()); + } + + void TearDown() override { + if (writer_) { + ASSERT_NO_THROW(writer_->Close()); + } + } + + SchemaDescriptor schema_; + std::shared_ptr properties_; + std::unique_ptr metadata_; + std::shared_ptr writer_; + std::shared_ptr arrow_properties_ = + default_arrow_writer_properties(); + ArrowWriteContext context_{::arrow::default_memory_pool(), arrow_properties_.get()}; +}; + +// Check that incompatible writes are rejected and valid shortcut writes still work +TEST_F(TestFixedSizeListWriter, RejectsMixedWritingMethods) { + MakeWriter(Type::INT32); + auto values = ::arrow::ArrayFromJSON(::arrow::int32(), "[1, 2, 3, 4, 5, 6]"); + context_.fixed_size_list_size = 3; + ASSERT_OK(writer_->WriteArrow(nullptr, nullptr, 6, *values, &context_, true)); + EXPECT_EQ(2, writer_->rows_written()); + EXPECT_EQ(0, writer_->estimated_buffered_def_level_bytes()); + EXPECT_EQ(0, writer_->estimated_buffered_rep_level_bytes()); + + auto typed = std::static_pointer_cast(writer_); + const int16_t definitions[] = {3, 2, 3}; + const int16_t repetitions[] = {0, 1, 1}; + const int32_t packed[] = {7, 9}; + const int32_t spaced[] = {7, 0, 9}; + const uint8_t validity = 0x05; + EXPECT_THROW(typed->WriteBatch(3, definitions, repetitions, packed), ParquetException); + EXPECT_THROW(typed->WriteBatchSpaced(3, definitions, repetitions, &validity, 0, spaced), + ParquetException); + auto packed_array = ::arrow::ArrayFromJSON(::arrow::int32(), "[7, 9]"); + context_.fixed_size_list_size = 0; + EXPECT_TRUE( + writer_->WriteArrow(definitions, repetitions, 3, *packed_array, &context_, true) + .IsInvalid()); + context_.fixed_size_list_size = 3; + ASSERT_OK(writer_->WriteArrow(nullptr, nullptr, 6, *values, &context_, true)); + EXPECT_EQ(4, writer_->rows_written()); +} + +// Check that the numeric shortcut rejects strings when called directly +TEST_F(TestFixedSizeListWriter, RejectsUnsupportedType) { + MakeWriter(Type::BYTE_ARRAY); + auto values = ::arrow::ArrayFromJSON(::arrow::utf8(), R"(["a", "b", "c"])"); + context_.fixed_size_list_size = 3; + EXPECT_TRUE( + writer_->WriteArrow(nullptr, nullptr, 3, *values, &context_, true).IsInvalid()); + EXPECT_EQ(0, writer_->rows_written()); +} + +// Check that callers cannot supply both marker arrays and a list-size shortcut +TEST_F(TestFixedSizeListWriter, RejectsMarkersWithRecipe) { + MakeWriter(Type::INT32); + auto values = ::arrow::ArrayFromJSON(::arrow::int32(), "[1, 2, 3]"); + const int16_t definitions[] = {3, 3, 3}; + const int16_t repetitions[] = {0, 1, 1}; + context_.fixed_size_list_size = 3; + EXPECT_TRUE(writer_->WriteArrow(definitions, repetitions, 3, *values, &context_, true) + .IsInvalid()); + EXPECT_EQ(0, writer_->rows_written()); +} + void GenerateLevels(int min_repeat_factor, int max_repeat_factor, int max_level, std::vector& input_levels) { // for each repetition count up to max_repeat_factor diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index 89f2be42bff..1618df3aa94 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -1591,6 +1591,9 @@ struct ArrowWriteContext { MemoryPool* memory_pool; const ArrowWriterProperties* properties; + // Use a positive list size only when every list and value is present + int32_t fixed_size_list_size = 0; + // Buffer used for storing the data of an array converted to the physical type // as expected by parquet-cpp. std::shared_ptr data_buffer;