From 917bfe01ab97aca3a14aaf11433cf77f960f75f1 Mon Sep 17 00:00:00 2001 From: Mashrfee Aryan Date: Mon, 12 Oct 2026 14:05:42 -0500 Subject: [PATCH 1/7] Add Parquet benchmark for writing fixed-size lists Compare writing fixed-size lists of floats with writing the same values as a flat column. Test six list sizes with V1 and V2 pages. --- .../parquet/arrow/reader_writer_benchmark.cc | 53 +++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/cpp/src/parquet/arrow/reader_writer_benchmark.cc b/cpp/src/parquet/arrow/reader_writer_benchmark.cc index 2f288fd2eb0f..a1eaec7234ee 100644 --- a/cpp/src/parquet/arrow/reader_writer_benchmark.cc +++ b/cpp/src/parquet/arrow/reader_writer_benchmark.cc @@ -222,6 +222,59 @@ BENCHMARK_TEMPLATE2(BM_WriteColumn, true, DoubleType); BENCHMARK_TEMPLATE2(BM_WriteColumn, false, BooleanType); BENCHMARK_TEMPLATE2(BM_WriteColumn, true, BooleanType); +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)); +} + +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, From 909cd6566fd72e3b0893924cea8806d06bc993c5 Mon Sep 17 00:00:00 2001 From: Mashrfee Aryan Date: Tue, 20 Oct 2026 21:52:22 -0500 Subject: [PATCH 2/7] Encode repeated values using a count Add an encoder helper that writes the same value many times. Check that it produces the same bytes as writing values one at a time. --- cpp/src/arrow/util/rle_encoding_internal.h | 24 ++++++++ cpp/src/arrow/util/rle_encoding_test.cc | 69 ++++++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index 4bd08fd65306..d0ed97a0bda8 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 up to count copies of value and return the number 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,27 @@ 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) { + 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 16b702c23d49..bd6fb689d579 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,72 @@ TEST(RleBitPacked, GetBatchSpacedRoundtripUint64) { DoTestGetBatchSpacedRoundtrip(); } +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); + } + } + } + } +} + +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 From 16ab2436f2748ce0167ea7f8652b5b58bf8a0518 Mon Sep 17 00:00:00 2001 From: Mashrfee Aryan Date: Tue, 20 Oct 2026 21:52:23 -0500 Subject: [PATCH 3/7] Encode definition levels using a count Write a repeated definition level without first building an array. Add tests for matching bytes and invalid inputs. --- cpp/src/parquet/column_writer.cc | 10 ++++++++ cpp/src/parquet/column_writer.h | 3 +++ cpp/src/parquet/column_writer_test.cc | 35 +++++++++++++++++++++++++++ 3 files changed, 48 insertions(+) diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index 3b4a873b5192..4ee3b4b56d50 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -242,6 +242,16 @@ int LevelEncoder::Encode(int batch_size, const int16_t* levels) { return num_encoded; } +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; +} + // ---------------------------------------------------------------------- // PageWriter implementation diff --git a/cpp/src/parquet/column_writer.h b/cpp/src/parquet/column_writer.h index 03817a77f317..731cd0bcb82c 100644 --- a/cpp/src/parquet/column_writer.h +++ b/cpp/src/parquet/column_writer.h @@ -73,6 +73,9 @@ 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 a constant level without an input array. Requires RLE encoding. + int EncodeRepeated(int batch_size, int16_t level); + 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 69ec07c1b78f..5602483288f8 100644 --- a/cpp/src/parquet/column_writer_test.cc +++ b/cpp/src/parquet/column_writer_test.cc @@ -1245,6 +1245,41 @@ TEST(TestColumnWriter, RepeatedListsUpdateSpacedBug) { writer->Close(); } +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); + } + } + } +} + +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); +} + 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 From 1af67b6104b22c9f837888188f42d7cefec3e8d9 Mon Sep 17 00:00:00 2001 From: Mashrfee Aryan Date: Tue, 20 Oct 2026 21:52:24 -0500 Subject: [PATCH 4/7] Encode repetition levels using the list size Use the fixed list size to write the start and continuation markers. Check that the encoded bytes match the existing encoder. --- cpp/src/parquet/column_writer.cc | 25 +++++++++++++++++++ cpp/src/parquet/column_writer.h | 3 +++ cpp/src/parquet/column_writer_test.cc | 36 +++++++++++++++++++++++++++ 3 files changed, 64 insertions(+) diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index 4ee3b4b56d50..e5a72b12dd5e 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -252,6 +252,31 @@ int LevelEncoder::EncodeRepeated(int batch_size, int16_t level) { return encoded; } +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 diff --git a/cpp/src/parquet/column_writer.h b/cpp/src/parquet/column_writer.h index 731cd0bcb82c..3eafc5f5a8cb 100644 --- a/cpp/src/parquet/column_writer.h +++ b/cpp/src/parquet/column_writer.h @@ -76,6 +76,9 @@ class PARQUET_EXPORT LevelEncoder { // Encode a constant level without an input array. Requires RLE encoding. int EncodeRepeated(int batch_size, int16_t level); + // Encode 0 followed by list_size - 1 ones, starting at a list boundary. + 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 5602483288f8..ca9604975e6b 100644 --- a/cpp/src/parquet/column_writer_test.cc +++ b/cpp/src/parquet/column_writer_test.cc @@ -1280,6 +1280,42 @@ TEST(TestLevelEncoder, EncodeRepeatedRejectsInvalidInput) { EXPECT_THROW(encoder.EncodeRepeated(1, 1), ParquetException); } +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); + } + } +} + +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); +} + 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 From 5d2210911ca965ddcf932bf1af52307e7b0f7e44 Mon Sep 17 00:00:00 2001 From: Mashrfee Aryan Date: Tue, 20 Oct 2026 21:52:25 -0500 Subject: [PATCH 5/7] Use the fast path to write fixed-size lists Skip temporary marker arrays for supported fixed-size lists without nulls. Keep the existing path for other inputs. Test matching output, lower memory use, and fallback cases. --- .../parquet/arrow/arrow_reader_writer_test.cc | 222 ++++++++++++++++++ cpp/src/parquet/arrow/writer.cc | 88 ++++++- cpp/src/parquet/column_writer.cc | 156 ++++++++++-- cpp/src/parquet/column_writer_test.cc | 81 +++++++ cpp/src/parquet/properties.h | 3 + 5 files changed, 519 insertions(+), 31 deletions(-) diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index d858962d4728..e7032ccf5c4f 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,227 @@ TEST(ArrowReadWrite, EmptyListView) { ASSERT_EQ(0, list_view.value_sizes()->size()); } +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)); +} + +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(); +} + +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)); + } + } + } + } + } + } +} + +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>()); +} + +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(); +} + +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); +} + +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/writer.cc b/cpp/src/parquet/arrow/writer.cc index e0fbe308219c..a40d1d22e887 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -135,6 +135,41 @@ Status GetSchemaMetadata(const ::arrow::Schema& schema, ::arrow::MemoryPool* poo return Status::OK(); } +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 +200,29 @@ 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); + 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 +299,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 +342,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 +351,23 @@ class ArrowColumnWriterV2 { } values_written += chunk_write_size; } - return std::make_unique(std::move(builders), column_index, - leaf_count, row_group_writer); + 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 +375,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 e5a72b12dd5e..6073858e7e2d 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -874,7 +874,8 @@ class ColumnWriterImpl { // RLE encode the src_buffer into dest_buffer and return the encoded size 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(); @@ -896,6 +897,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 @@ -970,25 +973,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(); @@ -1013,7 +1026,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 = @@ -1270,6 +1284,39 @@ inline void DoInBatchesRepeated(const int16_t* def_levels, const int16_t* rep_le } } +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 { + 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, @@ -1364,6 +1411,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 " @@ -1398,15 +1448,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 " @@ -1459,6 +1519,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 = @@ -1704,6 +1802,20 @@ 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) { + 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_test.cc b/cpp/src/parquet/column_writer_test.cc index ca9604975e6b..4ee77317fd95 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" @@ -1316,6 +1319,84 @@ TEST(TestLevelEncoder, EncodeFixedSizeListRejectsInvalidInput) { 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()}; +}; + +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()); +} + +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()); +} + +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 89f2be42bffd..932fd350caf2 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; + // Positive only for a column of complete, non-null fixed-size lists. + 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; From bd0e5447cd1db84aafdce5620ea13065c998649a Mon Sep 17 00:00:00 2001 From: Mashrfee Aryan Date: Tue, 27 Oct 2026 11:57:11 -0500 Subject: [PATCH 6/7] Fix benchmark formatting --- cpp/src/parquet/arrow/reader_writer_benchmark.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/parquet/arrow/reader_writer_benchmark.cc b/cpp/src/parquet/arrow/reader_writer_benchmark.cc index a1eaec7234ee..0fef36e7565a 100644 --- a/cpp/src/parquet/arrow/reader_writer_benchmark.cc +++ b/cpp/src/parquet/arrow/reader_writer_benchmark.cc @@ -255,7 +255,7 @@ static void BM_WriteFloatLayout(::benchmark::State& state) { while (state.KeepRunning()) { auto output = CreateOutputStream(); EXIT_NOT_OK(WriteTable(*table, ::arrow::default_memory_pool(), output, - table->num_rows(), properties)); + table->num_rows(), properties)); } const int64_t total_values = state.iterations() * num_values; From 94a6e719531e4180cf3b0b8bce42dee1d5d4342e Mon Sep 17 00:00:00 2001 From: Mashrfee Aryan Date: Wed, 28 Oct 2026 08:02:42 -0500 Subject: [PATCH 7/7] Explain tests and writer helpers --- cpp/src/arrow/util/rle_encoding_internal.h | 3 ++- cpp/src/arrow/util/rle_encoding_test.cc | 2 ++ cpp/src/parquet/arrow/arrow_reader_writer_test.cc | 7 +++++++ cpp/src/parquet/arrow/reader_writer_benchmark.cc | 2 ++ cpp/src/parquet/arrow/writer.cc | 3 +++ cpp/src/parquet/column_writer.cc | 7 ++++++- cpp/src/parquet/column_writer.h | 4 ++-- cpp/src/parquet/column_writer_test.cc | 7 +++++++ cpp/src/parquet/properties.h | 2 +- 9 files changed, 32 insertions(+), 5 deletions(-) diff --git a/cpp/src/arrow/util/rle_encoding_internal.h b/cpp/src/arrow/util/rle_encoding_internal.h index d0ed97a0bda8..08b3a8eaddfc 100644 --- a/cpp/src/arrow/util/rle_encoding_internal.h +++ b/cpp/src/arrow/util/rle_encoding_internal.h @@ -710,7 +710,7 @@ class RleBitPackedEncoder { /// This value must be representable with bit_width_ bits. bool Put(uint64_t value); - /// Encode up to count copies of value and return the number accepted. + /// 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. @@ -1536,6 +1536,7 @@ inline int RleBitPackedEncoder::PutRepeated(uint64_t value, int count) { 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; diff --git a/cpp/src/arrow/util/rle_encoding_test.cc b/cpp/src/arrow/util/rle_encoding_test.cc index bd6fb689d579..16f7bd9b9b6d 100644 --- a/cpp/src/arrow/util/rle_encoding_test.cc +++ b/cpp/src/arrow/util/rle_encoding_test.cc @@ -1454,6 +1454,7 @@ 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 = @@ -1494,6 +1495,7 @@ TEST(RleBitPacked, PutRepeatedMatchesPut) { } } +// 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)); diff --git a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc index e7032ccf5c4f..a60a82179873 100644 --- a/cpp/src/parquet/arrow/arrow_reader_writer_test.cc +++ b/cpp/src/parquet/arrow/arrow_reader_writer_test.cc @@ -3532,6 +3532,7 @@ 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, @@ -3576,6 +3577,7 @@ MakeFixedListTestTables(int32_t list_size, bool root_nullable, bool item_nullabl 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) { @@ -3608,6 +3610,7 @@ static Result> WriteFixedListTestTable( return sink->Finish(); } +// Compare saved bytes and read-back data across different writer settings template static void CheckFixedSizeListFastPath() { struct Settings { @@ -3678,6 +3681,7 @@ static void CheckFixedSizeListFastPath() { } } +// 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>()); @@ -3685,6 +3689,7 @@ TEST(ArrowReadWrite, FixedSizeListFastPathMatchesList) { 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()); { @@ -3705,6 +3710,7 @@ static Result MeasureFixedListWriterMemory(const Table& table) { 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; @@ -3719,6 +3725,7 @@ TEST(ArrowReadWrite, FixedSizeListFastPathUsesLessMemory) { 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]}])"}, diff --git a/cpp/src/parquet/arrow/reader_writer_benchmark.cc b/cpp/src/parquet/arrow/reader_writer_benchmark.cc index 0fef36e7565a..0471f90ea40b 100644 --- a/cpp/src/parquet/arrow/reader_writer_benchmark.cc +++ b/cpp/src/parquet/arrow/reader_writer_benchmark.cc @@ -222,6 +222,7 @@ 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)); @@ -263,6 +264,7 @@ static void BM_WriteFloatLayout(::benchmark::State& state) { 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}) { diff --git a/cpp/src/parquet/arrow/writer.cc b/cpp/src/parquet/arrow/writer.cc index a40d1d22e887..90380f30ef82 100644 --- a/cpp/src/parquet/arrow/writer.cc +++ b/cpp/src/parquet/arrow/writer.cc @@ -135,6 +135,7 @@ 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; @@ -215,6 +216,7 @@ class ArrowColumnWriterV2 { 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(), @@ -351,6 +353,7 @@ class ArrowColumnWriterV2 { } values_written += chunk_write_size; } + // 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, diff --git a/cpp/src/parquet/column_writer.cc b/cpp/src/parquet/column_writer.cc index 6073858e7e2d..89b9bfd892fd 100644 --- a/cpp/src/parquet/column_writer.cc +++ b/cpp/src/parquet/column_writer.cc @@ -242,6 +242,7 @@ 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_)) { @@ -252,6 +253,7 @@ int LevelEncoder::EncodeRepeated(int batch_size, int16_t level) { 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) { @@ -872,7 +874,7 @@ 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, bool is_repetition = false); @@ -1284,6 +1286,7 @@ 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, @@ -1309,6 +1312,7 @@ inline void DoInBatchesFixedSizeList(int64_t num_levels, int32_t 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); @@ -1803,6 +1807,7 @@ 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_) { diff --git a/cpp/src/parquet/column_writer.h b/cpp/src/parquet/column_writer.h index 3eafc5f5a8cb..86a3a56596a3 100644 --- a/cpp/src/parquet/column_writer.h +++ b/cpp/src/parquet/column_writer.h @@ -73,10 +73,10 @@ 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 a constant level without an input array. Requires RLE encoding. + // Encode one marker many times without an input array using RLE int EncodeRepeated(int batch_size, int16_t level); - // Encode 0 followed by list_size - 1 ones, starting at a list boundary. + // 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() { diff --git a/cpp/src/parquet/column_writer_test.cc b/cpp/src/parquet/column_writer_test.cc index 4ee77317fd95..4192143e36c5 100644 --- a/cpp/src/parquet/column_writer_test.cc +++ b/cpp/src/parquet/column_writer_test.cc @@ -1248,6 +1248,7 @@ 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}) { @@ -1270,6 +1271,7 @@ TEST(TestLevelEncoder, EncodeRepeatedMatchesEncode) { } } +// Check that invalid marker values, counts and encoder settings are rejected TEST(TestLevelEncoder, EncodeRepeatedRejectsInvalidInput) { LevelEncoder encoder; EXPECT_THROW(encoder.EncodeRepeated(1, 1), ParquetException); @@ -1283,6 +1285,7 @@ TEST(TestLevelEncoder, EncodeRepeatedRejectsInvalidInput) { 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}) { @@ -1306,6 +1309,7 @@ TEST(TestLevelEncoder, EncodeFixedSizeListMatchesEncode) { } } +// Check that invalid list sizes, counts and encoder settings are rejected TEST(TestLevelEncoder, EncodeFixedSizeListRejectsInvalidInput) { LevelEncoder encoder; EXPECT_THROW(encoder.EncodeFixedSizeList(3, 3), ParquetException); @@ -1349,6 +1353,7 @@ class TestFixedSizeListWriter : public ::testing::Test { 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]"); @@ -1377,6 +1382,7 @@ TEST_F(TestFixedSizeListWriter, RejectsMixedWritingMethods) { 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"])"); @@ -1386,6 +1392,7 @@ TEST_F(TestFixedSizeListWriter, RejectsUnsupportedType) { 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]"); diff --git a/cpp/src/parquet/properties.h b/cpp/src/parquet/properties.h index 932fd350caf2..1618df3aa94a 100644 --- a/cpp/src/parquet/properties.h +++ b/cpp/src/parquet/properties.h @@ -1591,7 +1591,7 @@ struct ArrowWriteContext { MemoryPool* memory_pool; const ArrowWriterProperties* properties; - // Positive only for a column of complete, non-null fixed-size lists. + // 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