diff --git a/benchmark/parquet_format_benchmark.cpp b/benchmark/parquet_format_benchmark.cpp index e8d3e0026..3823ac3bd 100644 --- a/benchmark/parquet_format_benchmark.cpp +++ b/benchmark/parquet_format_benchmark.cpp @@ -51,6 +51,7 @@ #include "arrow/c/helpers.h" #include "arrow/util/bit_util.h" #include "benchmark/benchmark.h" +#include "fmt/format.h" #include "paimon/common/utils/arrow/arrow_input_stream_adapter.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" @@ -104,6 +105,10 @@ constexpr int32_t kWriteBatchSize = 1024; constexpr int64_t kPageSizeBytes = 64 * 1024; // Four row groups per read fixture, so row-group pruning and page pruning are both in play. constexpr int64_t kRowGroupLength = 25'000; +// Only StringFixture lowers this from arrow's 1MB default: at 1MB a 25K-row row group of distinct +// `value_` entries still fits, so no cardinality this file writes would ever make the writer +// fall back to plain and the passthrough gate would have nothing to decline. +constexpr int64_t kDictionaryPageSizeBytes = 64 * 1024; constexpr int64_t kStringCardinality = 1'000; // Few enough distinct values that arrow keeps the column dictionary-encoded for the whole file, // which is the shape the wide-schema case wants: per-column work small, per-batch cost visible. @@ -160,6 +165,11 @@ std::shared_ptr DecimalSchema(int32_t precision) { return arrow::schema({MakeField("amount", arrow::decimal128(precision, 4), 0)}); } +// One STRING column, so a dictionary case measures one encoder and nothing else. +std::shared_ptr StringSchema() { + return arrow::schema({MakeField("name", arrow::utf8(), 0)}); +} + std::shared_ptr DoubleSchema() { return arrow::schema({MakeField("value", arrow::float64(), 0)}); } @@ -265,7 +275,8 @@ Result> MakeDictionaryColumn( return array; } -// STRING values: binary-like, so arrow can write the indices directly. +// STRING values: binary-like, so arrow can write the indices directly. Every batch is built over +// the same alphabet, the shape a single input file produces. Result> MakeDictionaryStringColumn(int64_t num_rows, int64_t offset, int64_t cardinality) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr values, @@ -273,6 +284,29 @@ Result> MakeDictionaryStringColumn(int64_t num_row return MakeDictionaryColumn(values, num_rows, offset); } +// The same, except every batch carries its own dictionary - the shape a compaction rewrite +// produces, where each input file supplies one. A Parquet column chunk holds a single dictionary, +// so arrow keeps the first and falls back to plain encoding for the rest of the row group +// (column_writer.cc, `dictionary->Equals(*preserved_dictionary_)`). +// +// The decoded column is byte for byte what MakeDictionaryStringColumn produces: same values, same +// order, same widths, same cardinality. Only the dictionary object differs, so the delta between +// the two benchmarks is the cost of the fallback and nothing else. Generating a fresh alphabet per +// batch instead would change the data as well - different strings, different widths, a different +// global cardinality - and the comparison would measure all of that at once. +Result> MakeChangingDictionaryStringColumn(int64_t num_rows, + int64_t offset, + int64_t cardinality) { + // The alphabet is rotated by one position per batch and the indices are shifted the other way, + // which cancels: index `(offset + cardinality - rotation + i) % cardinality` into an alphabet + // whose entry `j` is `value_<(j + rotation) % cardinality>` is `value_<(offset + i) % + // cardinality>` either way. `+ cardinality` keeps the shifted offset non-negative. + const int64_t rotation = num_rows > 0 ? (offset / num_rows) % cardinality : 0; + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr values, + MakeStringColumn(cardinality, rotation, cardinality)); + return MakeDictionaryColumn(values, num_rows, offset + cardinality - rotation); +} + // INT32 values: is_base_binary_like excludes them, so arrow densifies before writing. The // dictionary holds the same `i * 7` values MakeInt32Column emits inline, so the two are directly // comparable. @@ -422,6 +456,18 @@ BatchFactory SingleColumnBatch(const ColumnFactory& make_column) { }; } +// The same, but the batch is typed by the column rather than by the schema, so it can carry an +// encoding the schema does not declare. That is the shape a compaction rewrite produces: the file +// writer is built from the table's logical schema while the reader forwards whatever encoding the +// input file already had, leaving the writer to recover it from the batch. +BatchFactory SingleEncodedColumnBatch(const ColumnFactory& make_column) { + return [make_column](const std::shared_ptr& schema, int64_t offset, + int64_t rows) -> Result> { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr column, make_column(rows, offset)); + return MakeStructArray({schema->field(0)->WithType(column->type())}, {column}); + }; +} + Result> MakeNullableFlatBatch( const std::shared_ptr& schema, int64_t offset, int64_t rows, int64_t null_pct) { PAIMON_ASSIGN_OR_RAISE(std::shared_ptr ids, MakeInt64Column(rows, offset)); @@ -766,6 +812,27 @@ const ReadFixture& DoubleFixture() { return ColumnFixture("double", DoubleSchema(), &MakeDoubleColumn); } +// A one-column STRING file at a chosen cardinality, written with a reduced dictionary page limit +// so both regimes the passthrough gate distinguishes are reachable within a 100K-row file. Under +// the limit every data page stays dictionary-encoded and the gate lets the column through; over +// it the writer emits the dictionary page it has and encodes the rest as plain, which is the case +// the gate has to decline. At kDictionaryPageSizeBytes a row group holds roughly 4K distinct +// `value_` entries before overflowing, so cardinality alone picks the regime. +const ReadFixture& StringFixture(int64_t cardinality) { + const std::string key = fmt::format("string_{}", cardinality); + return GetFixture(key, [key, cardinality] { + std::map options; + options[paimon::parquet::PARQUET_DICTIONARY_PAGE_SIZE] = + std::to_string(kDictionaryPageSizeBytes); + return std::make_unique( + key + ".parquet", StringSchema(), + SingleColumnBatch([cardinality](int64_t rows, int64_t offset) { + return MakeStringColumn(rows, offset, cardinality); + }), + options); + }); +} + // The same data with dictionary encoding off, giving the read side a plain baseline. const ReadFixture& PlainFlatFixture() { return GetFixture("flat_plain", [] { @@ -930,6 +997,43 @@ void BM_ParquetWrite_DictionaryString(::benchmark::State& state) { kRowsPerBatch, /*options=*/{}); } +// arg: dictionary cardinality. The shape the append compaction rewrite actually produces, and the +// one BM_ParquetWrite_DictionaryString does not cover: there the schema itself is a DictionaryType, +// here the writer is built from a plain STRING schema - as a rewrite builds it, from the table's +// logical schema - and the batch arrives dictionary-encoded anyway. The delta against +// BM_ParquetWrite_String at the same cardinality is what the passthrough buys on the write side, +// including the per-batch schema fixup that recovers the encoding from the batch layout. +// +// One alphabet for the whole file, so the writer keeps writing indices: the favourable half of the +// passthrough. BM_ParquetWrite_ChangingDictionaryStringIntoStringSchema is the other half, and a +// rewrite of several input files lands between the two. +void BM_ParquetWrite_DictionaryStringIntoStringSchema(::benchmark::State& state) { + const int64_t cardinality = state.range(0); + RunWriteBenchmark(state, StringSchema(), + SingleEncodedColumnBatch([cardinality](int64_t rows, int64_t offset) { + return MakeDictionaryStringColumn(rows, offset, cardinality); + }), + kRowsPerBatch, /*options=*/{}, kDefaultCompression); +} + +// arg: dictionary cardinality. The same write, except every batch brings its own dictionary, as +// the input files of a compaction do, which is the cost of forwarding an encoding the writer +// cannot reuse. All three of BM_ParquetWrite_String, +// BM_ParquetWrite_DictionaryStringIntoStringSchema and this one write the same logical column at +// the same cardinality, so the triple reads directly: the middle one against the first is what the +// passthrough buys, this one against the middle is what the fallback costs in time, and this one +// against the first is the output size a rewrite that materialized and rebuilt would have produced. +// Row groups here are cut the way they are in production - by size and by the writer's memory +// limit, not at batch boundaries. +void BM_ParquetWrite_ChangingDictionaryStringIntoStringSchema(::benchmark::State& state) { + const int64_t cardinality = state.range(0); + RunWriteBenchmark(state, StringSchema(), + SingleEncodedColumnBatch([cardinality](int64_t rows, int64_t offset) { + return MakeChangingDictionaryStringColumn(rows, offset, cardinality); + }), + kRowsPerBatch, /*options=*/{}, kDefaultCompression); +} + // The same axis on an INTEGER dictionary, which arrow cannot direct-write - is_base_binary_like // excludes int32, so it densifies first. Its baseline is BM_ParquetWrite_FlatInt32 at the same // cardinality, not the String case: only the flat INT32 control holds value, width and encoding @@ -1197,6 +1301,22 @@ void BM_ParquetRead_Encoding(::benchmark::State& state, bool enable_dictionary) /*selection_bitmap=*/std::nullopt, /*options=*/{}, kReadBatchSize); } +// args: string cardinality, and whether the parquet dictionary passthrough is on. With it on, a +// column the file stores dictionary-encoded end to end is handed back as a DictionaryArray instead +// of one materialized value per row, so the pair at a fixed cardinality is what the read half of +// the compaction rewrite saves. At a cardinality high enough that the writer fell back to plain, +// the gate declines and the two runs measure the same work - a divergence there means the gate +// stopped looking at the data page encodings and started trusting the dictionary page. +void BM_ParquetRead_DictionaryPassthrough(::benchmark::State& state) { + const int64_t cardinality = state.range(0); + const bool enable_passthrough = state.range(1) != 0; + std::map options; + options[paimon::parquet::PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH] = + enable_passthrough ? "true" : "false"; + RunReadBenchmark(state, StringFixture(cardinality), StringSchema(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, options, kReadBatchSize); +} + // arg: decimal precision, the read side of BM_ParquetWrite_Decimal. Precision picks the physical // type - INT32, INT64 or FIXED_LEN_BYTE_ARRAY, since ParquetWriterBuilder enables // store_decimal_as_integer - and the three take different paths back to Decimal128Array. @@ -1272,6 +1392,24 @@ BENCHMARK(BM_ParquetWrite_DictionaryString) ->Arg(10000) ->Unit(benchmark::kMillisecond) ->UseRealTime(); +// Same cardinality axis as BM_ParquetWrite_String and BM_ParquetWrite_StringNoDictionary, which +// are its baselines: the three have to line up point for point or the low/medium/high comparison +// cannot be made. +BENCHMARK(BM_ParquetWrite_DictionaryStringIntoStringSchema) + ->ArgName("cardinality") + ->Arg(10) + ->Arg(1000) + ->Arg(kRowsPerFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); +// Same axis again, so the three points can be read against the run above them. +BENCHMARK(BM_ParquetWrite_ChangingDictionaryStringIntoStringSchema) + ->ArgName("cardinality") + ->Arg(10) + ->Arg(1000) + ->Arg(kRowsPerFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); BENCHMARK(BM_ParquetWrite_DictionaryInt32) ->ArgName("cardinality") ->Arg(10) @@ -1402,6 +1540,19 @@ BENCHMARK_CAPTURE(BM_ParquetRead_Encoding, dictionary, true) BENCHMARK_CAPTURE(BM_ParquetRead_Encoding, plain, false) ->Unit(benchmark::kMillisecond) ->UseRealTime(); +// The same cardinality axis the write cases use, so the read and write halves of a rewrite can be +// added up at each point. At kRowsPerFile every value is distinct, which overflows +// kDictionaryPageSizeBytes and is the point where the gate has to decline. +BENCHMARK(BM_ParquetRead_DictionaryPassthrough) + ->ArgNames({"cardinality", "passthrough"}) + ->Args({10, 0}) + ->Args({10, 1}) + ->Args({1000, 0}) + ->Args({1000, 1}) + ->Args({kRowsPerFile, 0}) + ->Args({kRowsPerFile, 1}) + ->Unit(benchmark::kMillisecond) + ->UseRealTime(); BENCHMARK(BM_ParquetRead_Decimal) ->ArgName("precision") ->Arg(9) diff --git a/benchmark/parquet_format_benchmark_test.cpp b/benchmark/parquet_format_benchmark_test.cpp index e930449d0..0bf8f4a7e 100644 --- a/benchmark/parquet_format_benchmark_test.cpp +++ b/benchmark/parquet_format_benchmark_test.cpp @@ -99,6 +99,17 @@ class ParquetFormatBenchmarkTest : public ::testing::Test { const std::shared_ptr& batch, const std::string& compression, const std::map& extra_options = {}, int32_t batch_count = 1) { + return WriteBatches(path, schema, + std::vector>(batch_count, batch), + compression, extra_options); + } + + // The same, for a benchmark case whose batches are not identical - the changing-dictionary + // write hands the writer a different dictionary every time. + Status WriteBatches(const std::string& path, const std::shared_ptr& schema, + const std::vector>& batches, + const std::string& compression, + const std::map& extra_options = {}) { std::map options = extra_options; // emplace, not assignment: a caller that set its own row-group limit is testing that. options.emplace(PARQUET_WRITE_MAX_ROW_GROUP_LENGTH, std::to_string(kRowGroupLength)); @@ -107,7 +118,7 @@ class ParquetFormatBenchmarkTest : public ::testing::Test { fs_->Create(path, /*overwrite=*/true)); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr writer, writer_builder.Build(out, compression)); - for (int32_t i = 0; i < batch_count; ++i) { + for (const std::shared_ptr& batch : batches) { ArrowArray c_array; PAIMON_RETURN_NOT_OK_FROM_ARROW(arrow::ExportArray(*batch, &c_array)); PAIMON_RETURN_NOT_OK(writer->AddBatch(&c_array)); @@ -116,14 +127,18 @@ class ParquetFormatBenchmarkTest : public ::testing::Test { return out->Close(); } - // Concatenating `array` with itself `times` times, so a multi-batch write has an expected - // value to be compared against. + // The expected value of a multi-batch write: one chunk per batch, in write order. + static Result> Concat( + const std::vector>& chunks) { + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr concatenated, + arrow::Concatenate(chunks)); + return concatenated; + } + + // The same, when every batch carried the same array. static Result> Repeat(const std::shared_ptr& array, int32_t times) { - std::vector> chunks(times, array); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr repeated, - arrow::Concatenate(chunks)); - return repeated; + return Concat(std::vector>(times, array)); } static Result> MakeDictionary( @@ -133,6 +148,54 @@ class ParquetFormatBenchmarkTest : public ::testing::Test { return array; } + // `value_`, the value MakeStringColumn emits at `index` in the benchmark. + static std::string AlphabetValue(int64_t index) { + return "value_" + std::to_string(index); + } + + // Indices `(shift + i) % cardinality` over kRows, the index column every dictionary case here + // is built from - the shape MakeDictionaryColumn gives a benchmark batch. `shift` cancels the + // rotation MakeChangingDictionaryStringColumn applies to its alphabet. + static Result> MakeIndices(int64_t cardinality, + int64_t shift = 0) { + arrow::Int32Builder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(kRows)); + for (int64_t i = 0; i < kRows; ++i) { + builder.UnsafeAppend(static_cast((shift + i) % cardinality)); + } + std::shared_ptr indices; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&indices)); + return indices; + } + + // The alphabet whose entry `j` is `value_<(j + rotation) % cardinality>`: the same + // `cardinality` strings whatever the rotation, in a different order. Rotating it is how + // MakeChangingDictionaryStringColumn gives every batch its own dictionary without changing the + // data the batch decodes to. + static Result> MakeAlphabet(int64_t cardinality, + int64_t rotation = 0) { + arrow::StringBuilder builder; + for (int64_t i = 0; i < cardinality; ++i) { + PAIMON_RETURN_NOT_OK_FROM_ARROW( + builder.Append(AlphabetValue((rotation + i) % cardinality))); + } + std::shared_ptr alphabet; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&alphabet)); + return alphabet; + } + + // The flat column any of those batches has to decode back to, whatever the rotation. + static Result> MakeFlatAlphabetColumn(int64_t cardinality) { + arrow::StringBuilder builder; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Reserve(kRows)); + for (int64_t i = 0; i < kRows; ++i) { + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Append(AlphabetValue(i % cardinality))); + } + std::shared_ptr column; + PAIMON_RETURN_NOT_OK_FROM_ARROW(builder.Finish(&column)); + return column; + } + struct ReadResult { int64_t rows = 0; uint64_t row_groups_total = 0; @@ -147,14 +210,15 @@ class ParquetFormatBenchmarkTest : public ::testing::Test { Result Read(const std::string& path, const std::shared_ptr& read_schema, const std::shared_ptr& predicate = nullptr, - const std::optional& selection = std::nullopt) { + const std::optional& selection = std::nullopt, + const std::map& read_options = {}) { PAIMON_ASSIGN_OR_RAISE(FileStatus file_status, fs_->GetFileStatus(path)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr input, fs_->Open(path)); auto in_stream = std::make_shared(input, file_status.GetLen(), pool_); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, ParquetFileBatchReader::Create( - std::move(in_stream), /*options=*/{}, kBatchSize, + std::move(in_stream), read_options, kBatchSize, /*file_metadata=*/nullptr, /*storage_read_bytes=*/nullptr, pool_, /*hints=*/std::nullopt)); ArrowSchema c_schema; @@ -237,35 +301,25 @@ TEST_F(ParquetFormatBenchmarkTest, RegisteredCodecsWrite) { // values match the flat equivalent. TEST_F(ParquetFormatBenchmarkTest, DictionaryInputRoundTrip) { constexpr int64_t kCardinality = 8; - arrow::Int32Builder index_builder; - ASSERT_TRUE(index_builder.Reserve(kRows).ok()); - for (int64_t i = 0; i < kRows; ++i) { - index_builder.UnsafeAppend(static_cast(i % kCardinality)); - } - std::shared_ptr indices; - ASSERT_TRUE(index_builder.Finish(&indices).ok()); - - arrow::StringBuilder string_values; + ASSERT_OK_AND_ASSIGN(std::shared_ptr indices, MakeIndices(kCardinality)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr string_dict, MakeAlphabet(kCardinality)); + // The flat array the dictionary-encoded input has to decode back to. + ASSERT_OK_AND_ASSIGN(std::shared_ptr flat_string_column, + MakeFlatAlphabetColumn(kCardinality)); + + // The INT32 dictionary holds the same `i * 7` values MakeInt32Column emits inline, which is + // what makes the two benchmark cases comparable. arrow::Int32Builder int_values; + arrow::Int32Builder flat_ints; for (int64_t i = 0; i < kCardinality; ++i) { - ASSERT_TRUE(string_values.Append("value_" + std::to_string(i)).ok()); ASSERT_TRUE(int_values.Append(static_cast(i * 7)).ok()); } - std::shared_ptr string_dict; - std::shared_ptr int_dict; - ASSERT_TRUE(string_values.Finish(&string_dict).ok()); - ASSERT_TRUE(int_values.Finish(&int_dict).ok()); - - // The flat arrays the dictionary-encoded input has to decode back to. - arrow::StringBuilder flat_strings; - arrow::Int32Builder flat_ints; for (int64_t i = 0; i < kRows; ++i) { - ASSERT_TRUE(flat_strings.Append("value_" + std::to_string(i % kCardinality)).ok()); ASSERT_TRUE(flat_ints.Append(static_cast((i % kCardinality) * 7)).ok()); } - std::shared_ptr flat_string_column; + std::shared_ptr int_dict; std::shared_ptr flat_int_column; - ASSERT_TRUE(flat_strings.Finish(&flat_string_column).ok()); + ASSERT_TRUE(int_values.Finish(&int_dict).ok()); ASSERT_TRUE(flat_ints.Finish(&flat_int_column).ok()); struct Case { @@ -303,6 +357,94 @@ TEST_F(ParquetFormatBenchmarkTest, DictionaryInputRoundTrip) { } } +// BM_ParquetWrite_DictionaryStringIntoStringSchema and its Changing... counterpart are only worth +// running as a pair if the pair measures two different things: one dictionary for the whole file +// against one per batch, where a Parquet column chunk's single dictionary forces the writer to +// plain encoding partway through. If that stopped happening the two would quietly measure the same +// work, so both halves are pinned here. The column shapes are rebuilt rather than taken from the +// benchmark, which this target does not compile in: the premise is what is pinned, not the +// factories. +// +// The two cases write *identical data* - same values, same order, same widths - and differ only in +// whether each batch's dictionary is a rotation of the last. That is what makes the benchmark delta +// attributable to the fallback, so this test asserts both halves against one expected column. +// +// The fallback is asserted through the reader's passthrough gate, which is the definition of +// "every data page dictionary-encoded" and decides whether a compacted file can be forwarded again +// on the next round. Its writer side, counted in pages, is +// ParquetFormatWriterTest.TestWriteDictionaryChangingAcrossBatches. +TEST_F(ParquetFormatBenchmarkTest, ChangingDictionaryInputFallsBackToPlain) { + constexpr int64_t kCardinality = 8; + constexpr int32_t kBatches = 3; + // Deliberately not a divisor of kRows, so a row group straddles two batches. A limit that + // divided the batch size would give every row group exactly one dictionary, the changing case + // could never fall back, and this test would pass while asserting nothing. + constexpr int64_t kStraddlingRowGroupLength = 700; + + // Plain STRING, as a compaction rewrite builds it from the table's logical schema. The batches + // carry the encoding the schema does not declare, which the writer recovers from their layout. + std::shared_ptr write_schema = arrow::schema({MakeField("v", arrow::utf8(), 0)}); + // One expected column for both cases, which is the point: rotating a dictionary is not + // supposed to change what the file holds. + ASSERT_OK_AND_ASSIGN(std::shared_ptr expected_chunk, + MakeFlatAlphabetColumn(kCardinality)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr expected, Repeat(expected_chunk, kBatches)); + + struct Case { + const char* name; + // Whether each batch rotates the dictionary, which is the only difference between the two + // column factories the benchmark pair uses. + bool changing; + // Whether the reader's gate still accepts the file, which is exactly whether the writer + // kept the column dictionary-encoded end to end. + bool expect_dictionary; + }; + for (const Case& c : {Case{"shared", false, true}, Case{"changing", true, false}}) { + std::vector> batches; + for (int32_t batch = 0; batch < kBatches; ++batch) { + // Rotating the alphabet and shifting the indices back by the same amount cancels, so + // every batch decodes to `expected_chunk` while presenting a dictionary the writer has + // not seen before. + const int64_t rotation = c.changing ? batch % kCardinality : 0; + ASSERT_OK_AND_ASSIGN(std::shared_ptr indices, + MakeIndices(kCardinality, kCardinality - rotation)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr alphabet, + MakeAlphabet(kCardinality, rotation)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr column, + MakeDictionary(indices, alphabet)); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr encoded, + Wrap(arrow::schema({write_schema->field(0)->WithType(column->type())}), {column})); + batches.push_back(std::move(encoded)); + } + + const std::string path = PathOf(std::string("changing_dict_") + c.name + ".parquet"); + ASSERT_OK(WriteBatches( + path, write_schema, batches, "zstd", + {{PARQUET_WRITE_MAX_ROW_GROUP_LENGTH, std::to_string(kStraddlingRowGroupLength)}})) + << c.name; + + // Values first: the fallback trades encoding and size, never correctness. + ASSERT_OK_AND_ASSIGN(ReadResult result, Read(path, write_schema)); + EXPECT_EQ(kRows * kBatches, result.rows) << c.name; + ASSERT_TRUE(result.data) << c.name; + EXPECT_TRUE( + checked_pointer_cast(result.data)->field(0)->Equals(*expected)) + << c.name; + + // Then the encoding, read back through the gate that decides it. + ASSERT_OK_AND_ASSIGN( + ReadResult encoded_result, + Read(path, write_schema, /*predicate=*/nullptr, /*selection=*/std::nullopt, + {{PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH, "true"}})); + ASSERT_TRUE(encoded_result.data) << c.name; + std::shared_ptr column = + checked_pointer_cast(encoded_result.data)->field(0); + EXPECT_EQ(c.expect_dictionary, column->type_id() == arrow::Type::DICTIONARY) + << c.name << ": " << column->type()->ToString(); + } +} + // DECIMAL precision selects the Parquet physical type, and precision 38 is the only one that // reaches FIXED_LEN_BYTE_ARRAY. The benchmark sweeps all three on both sides, so all three have // to survive a round trip with their type intact. diff --git a/docs/source/examples/benchmark.rst b/docs/source/examples/benchmark.rst index b82114eea..01d49cfb9 100644 --- a/docs/source/examples/benchmark.rst +++ b/docs/source/examples/benchmark.rst @@ -134,6 +134,25 @@ encoded; ``BM_ParquetWrite_Dictionary*`` vary whether the *input array* already is, which is what decides whether Arrow can pass indices through to Parquet or has to materialize them first. +Three of those cases form one comparison, at the same cardinality and over the +same logical column: ``BM_ParquetWrite_String`` writes it flat, +``BM_ParquetWrite_DictionaryStringIntoStringSchema`` writes it as one dictionary +forwarded through the whole file, and +``BM_ParquetWrite_ChangingDictionaryStringIntoStringSchema`` gives every batch +its own dictionary - a rotation of the same values, with the indices shifted the +other way, so the data is unchanged and only the dictionary object differs. A +Parquet column chunk holds one dictionary, so the third case makes the writer +fall back to plain encoding partway through the row group. Reading the second +against the first is what forwarding an encoding buys; the third against the +second is what that fallback costs in time; the third against the first is the +output size a rewrite that materialized and rebuilt would have produced. + +These are format-writer microbenchmarks: they measure ``AddBatch`` against a +Parquet file, not a compaction. The compaction time, CPU and peak memory of +``parquet.read.enable-dictionary-passthrough`` on a real table have to be +measured on that table - see the "Dictionary Passthrough" section of +:doc:`../user_guide/compaction`. + Reader cases (``BM_ParquetRead_*``) cover full scan, single-column projection, predicate-filtered reads at varying selectivity with page-index filtering on and off, skip-heavy reads driven by a strided selection bitmap, null density, diff --git a/docs/source/user_guide/compaction.rst b/docs/source/user_guide/compaction.rst index 093313555..492e27fff 100644 --- a/docs/source/user_guide/compaction.rst +++ b/docs/source/user_guide/compaction.rst @@ -88,6 +88,79 @@ After compaction, if the last output file is still smaller than ``compaction.file-size``, it is placed back into the compaction queue for future merging. +Dictionary Passthrough +~~~~~~~~~~~~~~~~~~~~~~ +An append-only compaction rewrite copies rows into the new file without +inspecting any value, so a Parquet column that an input file already stores +dictionary-encoded is forwarded to the writer still encoded instead of being +expanded to one copy of the value per row and re-encoded. This saves the reader +materializing the values and the writer hashing them again; how much that is +worth depends on the column, and low-cardinality ``STRING`` columns benefit most. +Primary-key compaction merges rows and is not covered. + +This is **off by default** and is enabled per table by setting +``parquet.read.enable-dictionary-passthrough`` to ``true``. It is opt-in rather +than automatic because of the output file size trade-off described below, which +depends on the data and is not a win on every table. Measure before turning it +on, and compare output file size as well as compaction time. + +Note that this is a *read* option, so it applies to every read of the table and +not only to the compaction rewrite. The values are unchanged, but an eligible +column reaches the consumer as an Arrow ``DictionaryArray`` rather than one value +per row, and a consumer that reads columns through its own accessors has to +unwrap it. Only a consumer that forwards batches without inspecting values - +which is what the rewrite does - gains anything from the encoding. + +Once enabled, eligibility is decided per input file: a non-nested ``STRING`` +column is forwarded when its data pages are dictionary-encoded throughout every +row group of *that* file, so one input file can be read encoded while the next +one is read as ordinary values, and the writer takes both. A high-cardinality +column that started dictionary-encoded and fell back to plain encoding therefore +does not qualify, even though it still carries a dictionary page. ``BINARY`` is +not forwarded although Parquet stores it in the same physical type and +dictionary-encodes it the same way, because the value accessors cannot read a +``BINARY`` dictionary. The rewrite also overrides the option back to ``false`` +when the table writes a format other than Parquet, when +``parquet.enable-dictionary`` is ``false`` because the writer would only expand +the values again, or when variant/map shredding is configured because those +writers reshape each batch against a fixed physical schema. Setting the option +therefore never makes a rewrite fail; at worst it has no effect. + +When the option is vetoed, the rewrite also enforces the veto on every input +batch. A format reader can independently hand over dictionary-encoded columns +because of its own lazy-decoding setting, so the rewrite decodes those columns +before handing them to a writer that cannot accept dictionary arrays. + +If a file index is configured on a forwarded column, that column alone is +materialized so the index still sees its values; the other columns stay encoded. + +Trade-off +^^^^^^^^^ +Passthrough changes what the rewrite costs, not what it produces. The rewritten +data is identical either way, but the output file can be **larger**. + +A Parquet column chunk can carry only one dictionary. The writer keeps the first +dictionary a column presents in a row group and falls back to plain encoding for +the rest of that row group as soon as a different one arrives. Compaction merges +several input files, each with its own dictionary, and output row groups are cut +by ``parquet.block.size`` and by the writer's memory limit, which are not aligned +to input file boundaries - a boundary may coincide, but nothing arranges for it. +So a rewrite can keep the first input file's dictionary and write the rest of the +row group plain, where materializing and rebuilding would have hashed the values +into a single dictionary for that row group (up to +``parquet.dictionary.page.size``, past which the writer falls back to plain in +either case). Where that happens the output column also becomes ineligible for +passthrough on the next compaction round, since its data pages are then only +partly dictionary-encoded. + +How often it happens, and what it costs on a given table, is what the +measurement above is for. + +Passthrough is therefore worth enabling when the reduction in compaction CPU +matters more than output size - for example when the same dictionary recurs +across input files, when the columns are wide enough that materializing them +dominates, or when the output is short-lived. Leave it off otherwise. + Append-Only Table Compaction Options ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/src/paimon/common/reader/reader_utils_test.cpp b/src/paimon/common/reader/reader_utils_test.cpp index 96b5a4da0..04639d931 100644 --- a/src/paimon/common/reader/reader_utils_test.cpp +++ b/src/paimon/common/reader/reader_utils_test.cpp @@ -28,8 +28,10 @@ #include "arrow/api.h" #include "arrow/array/array_base.h" #include "arrow/c/abi.h" +#include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "gtest/gtest.h" +#include "paimon/common/utils/checked_cast.h" #include "paimon/status.h" #include "paimon/testing/utils/read_result_collector.h" #include "paimon/testing/utils/testharness.h" @@ -104,4 +106,49 @@ TEST(ReaderUtilsTest, TestApplyBitmapToReadBatch) { "except eof"); } +TEST(ReaderUtilsTest, TestApplyBitmapToReadBatchKeepsDictionaryEncoding) { + // A deletion vector on an append table routes the Parquet dictionary passthrough through here: + // ParquetFileBatchReader reports SupportPreciseBitmapSelection() == false, so RawFileSplitRead + // wraps it and the surviving rows are cut out by slicing and concatenating. The encoding + // survives that only because every slice shares one dictionary and arrow::Concatenate has a + // fast path for it; if it ever unified or densified instead, a compaction with deletion + // vectors would quietly stop forwarding the encoding the rewrite asked for. + auto dictionary_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + auto make_encoded = [&dictionary_type](const std::string& indices_json) { + return arrow::DictionaryArray::FromArrays( + dictionary_type, + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), indices_json) + .ValueOrDie(), + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b"])") + .ValueOrDie()) + .ValueOrDie(); + }; + std::shared_ptr ids = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 2, 3, 4]").ValueOrDie(); + auto src_array = arrow::StructArray::Make({make_encoded("[0, 1, 0, 1, 0]"), ids}, + std::vector{"s", "id"}) + .ValueOrDie(); + + ASSERT_OK_AND_ASSIGN(auto src_batch, ReadResultCollector::GetReadBatch(src_array)); + // Two disjoint runs, so the filter has to concatenate rather than hand back a single slice. + auto batch_with_bitmap = + std::make_pair(std::move(src_batch), RoaringBitmap32::From(std::vector{0, 1, 4})); + ASSERT_OK_AND_ASSIGN(auto result_batch, + ReaderUtils::ApplyBitmapToReadBatch(std::move(batch_with_bitmap), + arrow::default_memory_pool())); + // Imported directly rather than through ReadResultCollector::GetArray, which decodes + // dictionaries on the way out and would hide the very thing being asserted. + auto& [c_array, c_schema] = result_batch; + std::shared_ptr result = + arrow::ImportArray(c_array.get(), c_schema.get()).ValueOrDie(); + ASSERT_EQ(3, result->length()); + auto result_struct = checked_pointer_cast(result); + ASSERT_EQ(arrow::Type::DICTIONARY, result_struct->field(0)->type()->id()); + ASSERT_TRUE(result_struct->field(0)->Equals(*make_encoded("[0, 1, 0]"))) + << "actual=" << result_struct->field(0)->ToString(); + std::shared_ptr expected_ids = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 4]").ValueOrDie(); + ASSERT_TRUE(result_struct->field(1)->Equals(*expected_ids)); +} + } // namespace paimon::test diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index 1ba4c153b..45884f5bb 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -24,6 +24,8 @@ #include "arrow/array/concatenate.h" #include "arrow/array/util.h" #include "arrow/buffer.h" +#include "arrow/c/abi.h" +#include "arrow/compute/cast.h" #include "arrow/type_traits.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" @@ -33,11 +35,61 @@ #include "paimon/common/utils/arrow/vector_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/string_utils.h" +#include "paimon/core/casting/casting_utils.h" namespace paimon { namespace { +// Whether `type` is a dictionary this can carry across the C data interface unchanged. The index +// width is part of the test because nothing in a layout reveals it; see +// ArrowUtils::IsDictionaryLayoutRecoverableValueType(). +bool IsDictionaryLayoutRecoverable(const arrow::DataType& type) { + if (type.id() != arrow::Type::DICTIONARY) { + return false; + } + const auto& dictionary_type = checked_cast(type); + return dictionary_type.index_type()->id() == arrow::Type::INT32 && + ArrowUtils::IsDictionaryLayoutRecoverableValueType(*dictionary_type.value_type()); +} + +// Whether `type` is or contains a dictionary at any depth. +bool HasDictionary(const arrow::DataType& type) { + if (type.id() == arrow::Type::DICTIONARY) { + return true; + } + for (const std::shared_ptr& field : type.fields()) { + if (HasDictionary(*field->type())) { + return true; + } + } + return false; +} + +// Whether any descendant of `array` carries a dictionary that `type` does not declare. `array` +// itself is not examined; its caller has already handled the top level. +bool HasUndeclaredDictionaryChild(const std::shared_ptr& type, + const ::ArrowArray* array) { + if (array == nullptr || array->n_children != type->num_fields()) { + return false; + } + for (int64_t i = 0; i < array->n_children; ++i) { + const ::ArrowArray* child = array->children[i]; + if (child == nullptr) { + continue; + } + const std::shared_ptr& child_type = + type->field(static_cast(i))->type(); + if (child->dictionary != nullptr && child_type->id() != arrow::Type::DICTIONARY) { + return true; + } + if (HasUndeclaredDictionaryChild(child_type, child)) { + return true; + } + } + return false; +} + bool NeedsNormalization(const std::shared_ptr& data) { if (data->offset != 0) { return true; @@ -500,4 +552,119 @@ Result ArrowUtils::GetCompressionType(const std::strin return compression_type; } +// `is_binary_like()` is BINARY and STRING and nothing else. `LARGE_STRING` is left out even though +// it is binary-like: the ORC reader widens strings to `dictionary(int64(), large_utf8())` under +// lazy decoding, and a layout reports neither index nor offset width, so reading that back as +// `int32` indices over `int32` offsets would silently reinterpret both buffers instead of failing. +// +// This narrows what may be carried; it cannot verify what was. See +// ResolveDictionaryStructTypeFromLayout() for where the index width becomes a caller contract. +bool ArrowUtils::IsDictionaryLayoutRecoverableValueType(const arrow::DataType& type) { + return arrow::is_binary_like(type.id()); +} + +// Why the header calls the `int32` index width a contract rather than a check: the value-type +// check rejects `dictionary(int64(), large_utf8())`, the shape the ORC reader produces, but +// nothing here can tell `dictionary(int32(), utf8())` apart from `dictionary(int64(), utf8())`, +// and the second would be read as the first. +// +// So the contract binds the producer, not the callers: `ParquetFormatWriter::ResolveBatchSchema` +// and `DataFileWriterBase::AddFileIndexBatch` see only the layout. +// `AppendOnlyFileStoreWrite::CompactRewrite` is today's only production path that can hand over a +// batch whose dictionaries the schema does not declare, and it honours the contract by running +// FlattenUnresolvableDictionaries() first. Closing the hole instead of narrowing it needs the real +// `ArrowSchema` to reach the writer, which `FormatWriter::AddBatch(ArrowArray*)` drops. +Result> ArrowUtils::ResolveDictionaryStructTypeFromLayout( + const std::shared_ptr& logical_type, const ::ArrowArray* batch) { + if (batch == nullptr || logical_type->id() != arrow::Type::STRUCT || + batch->n_children != logical_type->num_fields()) { + // Leave the mismatch to the import, which reports it with its own diagnostics. + return logical_type; + } + arrow::FieldVector fields; + bool has_dictionary = false; + for (int32_t i = 0; i < logical_type->num_fields(); ++i) { + const std::shared_ptr& field = logical_type->field(i); + const ::ArrowArray* child = batch->children[i]; + if (child == nullptr || child->dictionary == nullptr) { + if (HasUndeclaredDictionaryChild(field->type(), child)) { + return Status::NotImplemented(fmt::format( + "column '{}' is dictionary-encoded below its top level, which the Arrow " + "import cannot describe without the producer's schema", + field->name())); + } + fields.push_back(field); + continue; + } + if (field->type()->id() == arrow::Type::DICTIONARY) { + // The caller already declares the column as a dictionary, so its type describes the + // batch and nothing has to be recovered from the layout. + fields.push_back(field); + continue; + } + if (!IsDictionaryLayoutRecoverableValueType(*field->type())) { + return Status::NotImplemented(fmt::format( + "dictionary-encoded column '{}' of type {} cannot be resolved from the layout of " + "an ArrowArray, which pins down neither the index nor the offset width", + field->name(), field->type()->ToString())); + } + has_dictionary = true; + fields.push_back(field->WithType(arrow::dictionary(arrow::int32(), field->type()))); + } + if (!has_dictionary) { + return logical_type; + } + return arrow::struct_(fields); +} + +Result> ArrowUtils::FlattenUnresolvableDictionaries( + const std::shared_ptr& batch, + const std::shared_ptr& logical_type, arrow::MemoryPool* pool, + bool preserve_layout_recoverable_dictionaries) { + const std::shared_ptr& batch_type = batch->type(); + if (logical_type->id() != arrow::Type::STRUCT || !HasDictionary(*batch_type)) { + return batch; + } + const auto& logical_struct_type = checked_cast(*logical_type); + std::shared_ptr data; + arrow::FieldVector fields = batch_type->fields(); + for (int32_t i = 0; i < batch_type->num_fields(); ++i) { + std::shared_ptr field = fields[i]; + if (!HasDictionary(*field->type())) { + continue; + } + // Surviving the export is not enough when the destination imports against the logical + // type: an undeclared dictionary child of any shape then fails on the buffer count. + if (preserve_layout_recoverable_dictionaries && + IsDictionaryLayoutRecoverable(*field->type())) { + continue; + } + std::shared_ptr logical_field = + logical_struct_type.GetFieldByName(field->name()); + if (logical_field == nullptr) { + // Nothing says what this column should decode to, so leave it for the import to + // report against its own schema. + continue; + } + if (data == nullptr) { + // Copy once, on the first column that has to be decoded: the parent keeps its offset, + // length and validity, and only the child data is swapped underneath it. + data = batch->data()->Copy(); + } + // Decode the whole child rather than the slice the parent exposes, so the replacement + // lines up with the offset and length the parent still carries. + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr decoded, + CastingUtils::Cast(arrow::MakeArray(data->child_data[i]), logical_field->type(), + arrow::compute::CastOptions::Safe(), pool)); + data->child_data[i] = decoded->data(); + fields[i] = field->WithType(logical_field->type()); + } + if (data == nullptr) { + return batch; + } + data->type = arrow::struct_(fields); + return checked_pointer_cast(arrow::MakeArray(data)); +} + } // namespace paimon diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h index 59db09815..f78135073 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -26,6 +26,8 @@ #include "arrow/util/type_fwd.h" #include "paimon/result.h" +struct ArrowArray; + namespace paimon { class PAIMON_EXPORT ArrowUtils { @@ -72,6 +74,66 @@ class PAIMON_EXPORT ArrowUtils { /// Handles "none" and empty string by mapping them to "uncompressed". static Result GetCompressionType(const std::string& compression); + /// Whether `dictionary(int32(), type)` survives the Arrow C data interface, which drops the + /// type and leaves only the layout behind. `utf8()` and `binary()` do; `large_utf8()` does not, + /// because the layout reports neither the index nor the offset width. The definition says why. + /// + /// This is what the writer can recognise on the other side of the interface, not what a reader + /// should hand over: a producer may narrow it further for reasons of its own, and + /// ParquetFileBatchReader does, forwarding STRING alone. + /// + /// @param type The column's value type, not its dictionary type. + /// @return True when `dictionary(int32(), type)` round-trips through an `ArrowArray`. + static bool IsDictionaryLayoutRecoverableValueType(const arrow::DataType& type); + + /// Recovers dictionary fields omitted from a batch's declared logical type by inspecting its + /// layout: `logical_type` with every top-level field whose matching child in `batch` carries a + /// dictionary replaced by `dictionary(int32(), field type)`, or `logical_type` itself when no + /// child is dictionary-encoded. + /// + /// The `int32` index width is assumed rather than inferred, so this is a contract on whoever + /// produces the batch, not a check the callers can rely on: the producer must either provide a + /// batch whose dictionaries all have `int32` indices, or run + /// FlattenUnresolvableDictionaries() while the type is still known. The definition spells out + /// what that buys and what it does not. + /// + /// A field that already carries a dictionary type is left alone: `logical_type` then comes + /// from a caller that declared the encoding up front and already describes the batch. + /// + /// @param logical_type The struct type the caller declares for the batch. Returned unchanged + /// when it is not a struct or its field count does not match `batch`, + /// leaving the mismatch to the import's own diagnostics. + /// @param batch Only its structure is inspected, never its data, and it is not consumed. + /// @return `logical_type` or a copy of it carrying the recovered dictionary fields, or + /// NotImplemented for a dictionary this cannot describe. + static Result> ResolveDictionaryStructTypeFromLayout( + const std::shared_ptr& logical_type, const ::ArrowArray* batch); + + /// Returns a copy of `batch` in which every top-level column that cannot be preserved for the + /// destination has been decoded to the type its field carries in `logical_type`. A layout- + /// recoverable column may stay dictionary-encoded, so one column that has to be decoded does + /// not cost the others their encoding, and a batch that needs no decoding is returned + /// unchanged. + /// + /// The counterpart of the restriction above: an encoding the destination cannot take has to be + /// decoded while the type is still known. + /// + /// @param batch The batch to decode, matched to `logical_type` by field name; a column with no + /// matching field is left alone. + /// @param logical_type The struct type the decoded columns are cast to. `batch` is returned + /// unchanged when it is not a struct. + /// @param pool Allocates the decoded columns. Only used when a column is actually decoded. + /// @param preserve_layout_recoverable_dictionaries Whether dictionaries recoverable through + /// ResolveDictionaryStructTypeFromLayout() may remain encoded. Pass false + /// to decode every dictionary not already declared by `logical_type`, + /// whatever its shape. + /// @return `batch` itself when nothing had to be decoded, otherwise a copy of it with the + /// offset, length and validity of the original and the decoded columns swapped in. + static Result> FlattenUnresolvableDictionaries( + const std::shared_ptr& batch, + const std::shared_ptr& logical_type, arrow::MemoryPool* pool, + bool preserve_layout_recoverable_dictionaries); + private: static Status InnerCheckNullabilityMatch(const std::shared_ptr& field, const std::shared_ptr& data); diff --git a/src/paimon/common/utils/arrow/arrow_utils_test.cpp b/src/paimon/common/utils/arrow/arrow_utils_test.cpp index 2aad598b5..bc95a6d68 100644 --- a/src/paimon/common/utils/arrow/arrow_utils_test.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils_test.cpp @@ -19,12 +19,18 @@ #include "paimon/common/utils/arrow/arrow_utils.h" +#include +#include +#include + #include "arrow/api.h" #include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" #include "arrow/ipc/api.h" #include "gtest/gtest.h" #include "paimon/common/types/data_field.h" #include "paimon/common/utils/checked_cast.h" +#include "paimon/status.h" #include "paimon/testing/utils/testharness.h" namespace paimon::test { @@ -921,4 +927,292 @@ TEST(ArrowUtilsTest, TestGetCompressionType) { } } +TEST(ArrowUtilsTest, TestResolveDictionaryStructTypeFromLayout) { + // The resolution never consumes the exported batch, so it is released here. + auto resolve = [](const std::shared_ptr& array, + const std::shared_ptr& logical_type) { + ArrowArray c_array; + ArrowArrayMarkReleased(&c_array); + EXPECT_TRUE(arrow::ExportArray(*array, &c_array).ok()); + Result> resolved = + ArrowUtils::ResolveDictionaryStructTypeFromLayout(logical_type, &c_array); + ArrowArrayRelease(&c_array); + return resolved; + }; + + auto dictionary_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + std::shared_ptr indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 2]").ValueOrDie(); + std::shared_ptr strings = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["x", "yy", "zzz"])") + .ValueOrDie(); + std::shared_ptr ints = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[1, 2, 3]").ValueOrDie(); + std::shared_ptr encoded_strings = + arrow::DictionaryArray::FromArrays(dictionary_type, indices, strings).ValueOrDie(); + auto logical_type = + arrow::struct_({arrow::field("s", arrow::utf8()), arrow::field("i", arrow::int32())}); + + { + // No dictionary: the very same type instance comes back. + auto batch = arrow::StructArray::Make({strings, ints}, std::vector{"s", "i"}) + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr resolved, + resolve(batch, logical_type)); + ASSERT_EQ(logical_type, resolved); + } + { + // int32 indices, the only encoding Arrow's Parquet reader produces. + auto batch = + arrow::StructArray::Make({encoded_strings, ints}, std::vector{"s", "i"}) + .ValueOrDie(); + ASSERT_OK_AND_ASSIGN(std::shared_ptr resolved, + resolve(batch, logical_type)); + ASSERT_TRUE(resolved->Equals(*arrow::struct_( + {arrow::field("s", dictionary_type), arrow::field("i", arrow::int32())}))); + } + { + // A caller that declared the dictionary up front already describes the batch, so its + // type is preserved even for a value type the layout-derived path would reject. + std::shared_ptr encoded_ints = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::int32()), + indices, ints) + .ValueOrDie(); + auto batch = arrow::StructArray::Make({encoded_strings, encoded_ints}, + std::vector{"s", "i"}) + .ValueOrDie(); + auto declared_type = + arrow::struct_({arrow::field("s", dictionary_type), + arrow::field("i", arrow::dictionary(arrow::int32(), arrow::int32()))}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr resolved, + resolve(batch, declared_type)); + ASSERT_EQ(declared_type, resolved); + } + { + // Nothing in the layout pins down the index width, so a dictionary Arrow would never + // produce is rejected instead of being reinterpreted. + std::shared_ptr encoded_ints = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::int32()), + indices, ints) + .ValueOrDie(); + auto batch = + arrow::StructArray::Make({strings, encoded_ints}, std::vector{"s", "i"}) + .ValueOrDie(); + Status status = resolve(batch, logical_type).status(); + ASSERT_TRUE(status.IsNotImplemented()) << status.ToString(); + } + { + auto nested = + arrow::StructArray::Make({encoded_strings}, std::vector{"s"}).ValueOrDie(); + auto batch = arrow::StructArray::Make({nested, ints}, std::vector{"n", "i"}) + .ValueOrDie(); + // Same for a dictionary hidden below the top level. + auto undeclared_type = + arrow::struct_({arrow::field("n", arrow::struct_({arrow::field("s", arrow::utf8())})), + arrow::field("i", arrow::int32())}); + Status status = resolve(batch, undeclared_type).status(); + ASSERT_TRUE(status.IsNotImplemented()) << status.ToString(); + + // Unless the type declares it there too, which keeps the rule uniform with the top level. + auto declared_type = + arrow::struct_({arrow::field("n", arrow::struct_({arrow::field("s", dictionary_type)})), + arrow::field("i", arrow::int32())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr resolved, + resolve(batch, declared_type)); + ASSERT_EQ(declared_type, resolved); + } + { + // A layout says nothing about offset width either, so large_utf8 is rejected too. The ORC + // reader widens strings to dictionary(int64, large_utf8) under lazy decoding, and both of + // its buffers would be misread if this guessed int32 the way it does for utf8. + std::shared_ptr large_strings = + arrow::ipc::internal::json::ArrayFromJSON(arrow::large_utf8(), R"(["x", "yy"])") + .ValueOrDie(); + std::shared_ptr large_indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), "[0, 1, 0]").ValueOrDie(); + auto large_dictionary_type = arrow::dictionary(arrow::int64(), arrow::large_utf8()); + std::shared_ptr encoded_large_strings = + arrow::DictionaryArray::FromArrays(large_dictionary_type, large_indices, large_strings) + .ValueOrDie(); + auto batch = arrow::StructArray::Make({encoded_large_strings, ints}, + std::vector{"s", "i"}) + .ValueOrDie(); + auto large_logical_type = arrow::struct_( + {arrow::field("s", arrow::large_utf8()), arrow::field("i", arrow::int32())}); + Status status = resolve(batch, large_logical_type).status(); + ASSERT_TRUE(status.IsNotImplemented()) << status.ToString(); + } +} + +TEST(ArrowUtilsTest, TestFlattenUnresolvableDictionaries) { + auto pool = arrow::default_memory_pool(); + auto layout_recoverable_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + // What the ORC reader hands over for a dictionary-encoded string column under lazy decoding. + auto layout_unrecoverable_type = arrow::dictionary(arrow::int64(), arrow::large_utf8()); + auto logical_type = + arrow::struct_({arrow::field("s", arrow::utf8()), arrow::field("o", arrow::utf8()), + arrow::field("i", arrow::int32())}); + + std::shared_ptr ints = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[1, 2, 3]").ValueOrDie(); + std::shared_ptr layout_recoverable = + arrow::DictionaryArray::FromArrays( + layout_recoverable_type, + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 0]").ValueOrDie(), + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b"])").ValueOrDie()) + .ValueOrDie(); + std::shared_ptr layout_unrecoverable = + arrow::DictionaryArray::FromArrays( + layout_unrecoverable_type, + arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), "[0, null, 1]").ValueOrDie(), + arrow::ipc::internal::json::ArrayFromJSON(arrow::large_utf8(), R"(["c", "d"])") + .ValueOrDie()) + .ValueOrDie(); + + { + // Only the column that would not survive the export is decoded; the one that would keeps + // its encoding, which is what makes this selective rather than an all-or-nothing flatten. + auto batch = checked_pointer_cast( + arrow::StructArray::Make({layout_recoverable, layout_unrecoverable, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries( + batch, logical_type, pool, + /*preserve_layout_recoverable_dictionaries=*/true)); + ASSERT_EQ(arrow::Type::DICTIONARY, flattened->field(0)->type()->id()); + ASSERT_TRUE(flattened->field(1)->type()->Equals(*arrow::utf8())); + ASSERT_EQ(arrow::Type::INT32, flattened->field(2)->type()->id()); + + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["c", null, "d"])") + .ValueOrDie(); + ASSERT_TRUE(flattened->field(1)->Equals(*expected)) + << "actual=" << flattened->field(1)->ToString(); + ASSERT_TRUE(flattened->field(0)->Equals(*layout_recoverable)); + } + { + // A dictionary below the top level is layout-unrecoverable too, so the whole column is + // decoded. + auto nested = + arrow::StructArray::Make({layout_unrecoverable}, std::vector{"o"}) + .ValueOrDie(); + auto batch = checked_pointer_cast( + arrow::StructArray::Make({nested, ints}, std::vector{"n", "i"}) + .ValueOrDie()); + auto nested_logical_type = + arrow::struct_({arrow::field("n", arrow::struct_({arrow::field("o", arrow::utf8())})), + arrow::field("i", arrow::int32())}); + ASSERT_OK_AND_ASSIGN(std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries( + batch, nested_logical_type, pool, + /*preserve_layout_recoverable_dictionaries=*/true)); + ASSERT_TRUE(flattened->field(0)->type()->Equals( + *arrow::struct_({arrow::field("o", arrow::utf8())}))) + << flattened->field(0)->type()->ToString(); + } + { + // The case the value type alone cannot rule out: `utf8` values behind `int64` indices. + // ResolveDictionaryStructTypeFromLayout() would accept the value type and then read the + // indices as `int32`, so the index width has to be caught here or not at all - one of the + // two reasons CompactRewrite runs this before exporting rather than relying on the + // writer's own rejection. A destination that resolves nothing is the other, below. + std::shared_ptr wide_indices = + arrow::DictionaryArray::FromArrays( + arrow::dictionary(arrow::int64(), arrow::utf8()), + arrow::ipc::internal::json::ArrayFromJSON(arrow::int64(), "[1, 0, 1]").ValueOrDie(), + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["e", "f"])") + .ValueOrDie()) + .ValueOrDie(); + auto batch = checked_pointer_cast( + arrow::StructArray::Make({layout_recoverable, wide_indices, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries( + batch, logical_type, pool, + /*preserve_layout_recoverable_dictionaries=*/true)); + ASSERT_TRUE(flattened->field(1)->type()->Equals(*arrow::utf8())) + << flattened->field(1)->type()->ToString(); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["f", "e", "f"])") + .ValueOrDie(); + ASSERT_TRUE(flattened->field(1)->Equals(*expected)) + << "actual=" << flattened->field(1)->ToString(); + // The int32-indexed neighbour is untouched, so catching one does not cost the other. + ASSERT_EQ(arrow::Type::DICTIONARY, flattened->field(0)->type()->id()); + } + { + // Nothing to do: the very same array comes back, so a rewrite that never sees a dictionary + // pays nothing for this. + auto batch = checked_pointer_cast( + arrow::StructArray::Make({layout_recoverable, layout_recoverable, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries( + batch, logical_type, pool, + /*preserve_layout_recoverable_dictionaries=*/true)); + ASSERT_EQ(batch, flattened); + } + { + // A sliced batch keeps its offset: only the child data is swapped underneath it, so the + // rows the parent exposes stay the ones it exposed before. + auto batch = checked_pointer_cast( + arrow::StructArray::Make({layout_recoverable, layout_unrecoverable, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + auto sliced = checked_pointer_cast(batch->Slice(1, 2)); + ASSERT_OK_AND_ASSIGN(std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries( + sliced, logical_type, pool, + /*preserve_layout_recoverable_dictionaries=*/true)); + ASSERT_EQ(2, flattened->length()); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"([null, "d"])").ValueOrDie(); + ASSERT_TRUE(flattened->field(1)->Equals(*expected)) + << "actual=" << flattened->field(1)->ToString(); + } + { + // The first dictionary models a format reader whose own lazy-decoding option emits + // dictionary(int32, utf8). Its layout is recoverable, so the selective path would preserve + // it; a writer that imports against the logical type still needs it decoded. The second + // dictionary models the wider shape emitted by the built-in ORC reader. + auto batch = checked_pointer_cast( + arrow::StructArray::Make({layout_recoverable, layout_unrecoverable, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries( + batch, logical_type, pool, + /*preserve_layout_recoverable_dictionaries=*/false)); + ASSERT_TRUE(flattened->field(0)->type()->Equals(*arrow::utf8())) + << flattened->field(0)->type()->ToString(); + ASSERT_TRUE(flattened->field(1)->type()->Equals(*arrow::utf8())) + << flattened->field(1)->type()->ToString(); + ASSERT_EQ(arrow::Type::INT32, flattened->field(2)->type()->id()); + + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "a"])") + .ValueOrDie(); + ASSERT_TRUE(flattened->field(0)->Equals(*expected)) + << "actual=" << flattened->field(0)->ToString(); + } + { + // The batch the selective path returns by identity is copied and decoded here instead: + // "nothing to do" is a statement about the destination, not about the batch alone. + auto batch = checked_pointer_cast( + arrow::StructArray::Make({layout_recoverable, layout_recoverable, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN(std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries( + batch, logical_type, pool, + /*preserve_layout_recoverable_dictionaries=*/false)); + ASSERT_NE(batch, flattened); + ASSERT_TRUE(flattened->field(0)->type()->Equals(*arrow::utf8())); + ASSERT_TRUE(flattened->field(1)->type()->Equals(*arrow::utf8())); + } +} + } // namespace paimon::test diff --git a/src/paimon/core/casting/casting_utils.h b/src/paimon/core/casting/casting_utils.h index d28ff3cc7..8d27f06c4 100644 --- a/src/paimon/core/casting/casting_utils.h +++ b/src/paimon/core/casting/casting_utils.h @@ -27,10 +27,11 @@ #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/field_type_utils.h" #include "paimon/data/decimal.h" +#include "paimon/defs.h" #include "paimon/predicate/literal.h" namespace paimon { -class CastingUtils { +class PAIMON_EXPORT CastingUtils { public: CastingUtils() = delete; ~CastingUtils() = delete; diff --git a/src/paimon/core/io/data_file_index_writer.cpp b/src/paimon/core/io/data_file_index_writer.cpp index 5c97bb3da..78e6c4260 100644 --- a/src/paimon/core/io/data_file_index_writer.cpp +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -20,18 +20,22 @@ #include "paimon/core/io/data_file_index_writer.h" #include +#include #include #include "arrow/api.h" #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "arrow/compute/cast.h" #include "fmt/format.h" #include "paimon/common/io/byte_array_output_stream.h" #include "paimon/common/io/memory_segment_output_stream.h" #include "paimon/common/table/special_fields.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/path_util.h" #include "paimon/common/utils/scope_guard.h" +#include "paimon/core/casting/casting_utils.h" #include "paimon/core/io/data_file_path_factory.h" #include "paimon/file_index/file_index_format.h" #include "paimon/file_index/file_index_writer.h" @@ -96,14 +100,41 @@ DataFileIndexWriter::DataFileIndexWriter(std::vector&& writers path_factory_(path_factory), pool_(pool) {} +DataFileIndexWriter::~DataFileIndexWriter() = default; + Status DataFileIndexWriter::AddBatch(const std::shared_ptr& logical_batch) { if (finished_) { return Status::Invalid("Data file index writer has already finished"); } + // One entry per indexed column, not per index: a column carrying both a bitmap and a bloom + // filter appears twice in `writers_` and would otherwise be materialized twice per batch. + // Keyed by field index, which fixes the target type too - every entry for a column takes its + // `field` from the same position of the logical schema. + std::unordered_map> decoded_columns; for (const IndexWriterEntry& entry : writers_) { - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - std::shared_ptr projected, - arrow::StructArray::Make({logical_batch->field(entry.field_index)}, {entry.field})); + std::shared_ptr column = logical_batch->field(entry.field_index); + if (column->type_id() == arrow::Type::DICTIONARY && + entry.field->type()->id() != arrow::Type::DICTIONARY) { + // Index writers read values position by position, so a column forwarded encoded by + // the parquet dictionary passthrough is materialized first. Only indexed columns pay + // for this; the rest reach the data file writer still encoded. Materializing a large + // string column is worth accounting for, hence the project pool rather than Arrow's. + auto cached = decoded_columns.find(entry.field_index); + if (cached != decoded_columns.end()) { + column = cached->second; + } else { + if (arrow_pool_ == nullptr) { + arrow_pool_ = GetArrowPool(pool_); + } + PAIMON_ASSIGN_OR_RAISE( + column, + CastingUtils::Cast(column, entry.field->type(), + arrow::compute::CastOptions::Safe(), arrow_pool_.get())); + decoded_columns.emplace(entry.field_index, column); + } + } + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr projected, + arrow::StructArray::Make({column}, {entry.field})); ::ArrowArray c_array; ArrowArrayMarkReleased(&c_array); ScopeGuard array_guard([&c_array]() { ArrowArrayRelease(&c_array); }); diff --git a/src/paimon/core/io/data_file_index_writer.h b/src/paimon/core/io/data_file_index_writer.h index 883b719fc..ed70b4c5f 100644 --- a/src/paimon/core/io/data_file_index_writer.h +++ b/src/paimon/core/io/data_file_index_writer.h @@ -31,6 +31,7 @@ namespace arrow { class Field; +class MemoryPool; class Schema; class StructArray; } // namespace arrow @@ -57,6 +58,9 @@ class DataFileIndexWriter { const std::shared_ptr& path_factory, const std::shared_ptr& pool); + /// Out of line so `arrow_pool_` only needs a forward declaration here. + ~DataFileIndexWriter(); + Status AddBatch(const std::shared_ptr& logical_batch); /// Finalizes and publishes all configured indexes. This is a terminal, one-shot operation. @@ -88,6 +92,13 @@ class DataFileIndexWriter { Result> SerializeContainer(); Status WriteExternal(const std::string& path, const std::shared_ptr& bytes); + /// Decodes the dictionary-encoded columns the parquet passthrough forwards, and nothing else. + /// Built on first use, since most batches decode nothing. + /// + /// Declared before `writers_` on purpose: buffers allocated through the adaptor keep a raw + /// pointer to it, so it has to be destroyed last, after anything an index writer may still + /// hold. See ArrowMemPoolAdaptor in common/utils/arrow/mem_utils.cpp. + std::unique_ptr arrow_pool_; std::vector writers_; int64_t in_manifest_threshold_; std::shared_ptr file_system_; diff --git a/src/paimon/core/io/data_file_index_writer_test.cpp b/src/paimon/core/io/data_file_index_writer_test.cpp index 1e2c9593f..9837b5563 100644 --- a/src/paimon/core/io/data_file_index_writer_test.cpp +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -23,7 +23,10 @@ #include #include #include +#include +#include "arrow/array/array_dict.h" +#include "arrow/array/array_nested.h" #include "arrow/c/bridge.h" #include "arrow/ipc/json_simple.h" #include "arrow/type.h" @@ -169,6 +172,58 @@ TEST_F(DataFileIndexWriterTest, TestBitmapAndRangeBitmapEmbeddedRoundTrip) { ASSERT_EQ("{2,3}", greater_result->ToString()); } +TEST_F(DataFileIndexWriterTest, TestDictionaryEncodedIndexedColumnRoundTrip) { + // The parquet dictionary passthrough hands compaction batches over still encoded, and the + // bitmap index only sees the right values if the indexed column is decoded first. + schema_ = + arrow::schema({arrow::field("f0", arrow::utf8()), arrow::field("f1", arrow::int32())}); + ASSERT_OK_AND_ASSIGN(auto writer, + CreateWriter({{"file-index.bitmap.columns", "f0"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}})); + + auto encoded_batch = [](const std::string& indices_json, const std::string& dictionary_json, + const std::string& ints_json) { + std::shared_ptr indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), indices_json).ValueOrDie(); + std::shared_ptr dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), dictionary_json).ValueOrDie(); + std::shared_ptr encoded = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::utf8()), + indices, dictionary) + .ValueOrDie(); + std::shared_ptr values = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), ints_json).ValueOrDie(); + return checked_pointer_cast( + arrow::StructArray::Make({encoded, values}, std::vector{"f0", "f1"}) + .ValueOrDie()); + }; + + // Two batches with different dictionaries, which is what a rewrite of several input files + // hands over: each has to be decoded against its own alphabet, and the pool those decoded + // columns come from outlives the call that built it. + ASSERT_OK( + writer->AddBatch(encoded_batch("[0, 1, 0, 2]", R"(["a", "b", "c"])", "[10, 20, 30, 40]"))); + ASSERT_OK(writer->AddBatch(encoded_batch("[1, 0]", R"(["b", "d"])", "[50, 60]"))); + + ASSERT_OK_AND_ASSIGN(FileIndexWriteResult result, writer->Finish("unused.orc")); + ASSERT_TRUE(result.embedded_index); + ASSERT_OK_AND_ASSIGN(auto reader, CreateReader(result.embedded_index)); + + ASSERT_OK_AND_ASSIGN(auto bitmap_readers, ReadColumn(reader.get(), "f0")); + ASSERT_EQ(1, bitmap_readers.size()); + ASSERT_OK_AND_ASSIGN(auto equal_result, + bitmap_readers[0]->VisitEqual(Literal(FieldType::STRING, "a", 1))); + ASSERT_EQ("{0,2}", equal_result->ToString()); + // Row 1 comes from the first dictionary and row 5 from the second, so a decode that reused the + // wrong alphabet would land somewhere else. + ASSERT_OK_AND_ASSIGN(auto shared_value_result, + bitmap_readers[0]->VisitEqual(Literal(FieldType::STRING, "b", 1))); + ASSERT_EQ("{1,5}", shared_value_result->ToString()); + ASSERT_OK_AND_ASSIGN(auto second_batch_result, + bitmap_readers[0]->VisitEqual(Literal(FieldType::STRING, "d", 1))); + ASSERT_EQ("{4}", second_batch_result->ToString()); +} + TEST_F(DataFileIndexWriterTest, TestExternalIndexAndAbortCleanup) { ASSERT_OK_AND_ASSIGN(auto writer, CreateWriter({{"file-index.bitmap.columns", "f0"}, diff --git a/src/paimon/core/io/data_file_writer_base.h b/src/paimon/core/io/data_file_writer_base.h index ccea898a7..9b3d6e71d 100644 --- a/src/paimon/core/io/data_file_writer_base.h +++ b/src/paimon/core/io/data_file_writer_base.h @@ -25,7 +25,9 @@ #include #include "arrow/c/bridge.h" +#include "arrow/c/helpers.h" #include "arrow/type.h" +#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/core/io/data_file_index_writer.h" #include "paimon/core/io/data_file_meta.h" @@ -127,8 +129,20 @@ class DataFileWriterBase : public SingleFileWriter> batch_type = + ArrowUtils::ResolveDictionaryStructTypeFromLayout(logical_type_, batch); + if (!batch_type.ok()) { + // Every other exit from here has already handed `batch` to ImportArray, which consumes + // it whether it succeeds or not. Keep that contract on the one path that returns + // before the import runs, or a caller holding the array only in a local would leak it. + ArrowArrayRelease(batch); + return batch_type.status(); + } PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr logical_array, - arrow::ImportArray(batch, logical_type_)); + arrow::ImportArray(batch, batch_type.value())); std::shared_ptr logical_batch = checked_pointer_cast(logical_array); PAIMON_RETURN_NOT_OK(file_index_writer_->AddBatch(logical_batch)); diff --git a/src/paimon/core/operation/append_only_file_store_write.cpp b/src/paimon/core/operation/append_only_file_store_write.cpp index f660093a2..7448a6d96 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -30,7 +30,9 @@ #include "paimon/common/data/shredding/shredding_write_plan_factories.h" #include "paimon/common/table/special_fields.h" #include "paimon/common/utils/arrow/arrow_utils.h" +#include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" +#include "paimon/common/utils/options_utils.h" #include "paimon/core/append/append_only_writer.h" #include "paimon/core/append/bucketed_append_compact_manager.h" #include "paimon/core/compact/noop_compact_manager.h" @@ -68,6 +70,22 @@ class DataFilePathFactory; class MemoryPool; class SchemaManager; +namespace { + +// Spelled out rather than taken from `paimon/format/parquet/parquet_format_defs.h`: the format +// layer is pluggable, and an engine that supplies its own Parquet implementation need not export +// those symbols. The two defaults restate what the Parquet layer resolves these options to when +// the table does not set them - ParquetWriterBuilder for the first, ParquetFileBatchReader for the +// second. +constexpr char kParquetFormat[] = "parquet"; +constexpr char kParquetEnableDictionary[] = "parquet.enable-dictionary"; +constexpr char kParquetReadEnableDictionaryPassthrough[] = + "parquet.read.enable-dictionary-passthrough"; +constexpr bool kDefaultParquetEnableDictionary = true; +constexpr bool kDefaultParquetReadEnableDictionaryPassthrough = false; + +} // namespace + AppendOnlyFileStoreWrite::AppendOnlyFileStoreWrite( const std::shared_ptr& file_store_path_factory, const std::shared_ptr& snapshot_manager, @@ -146,13 +164,25 @@ Result>> AppendOnlyFileStoreWrite::Com return std::vector>{}; } - PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, - CreateFilesReader(partition, bucket, dv_factory, to_compact)); + // Resolved once: the reader and the writer have to agree on whether this rewrite stays a + // passthrough, and selecting the plan twice would let them drift apart. + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr plan_factory, + ShreddingWritePlanFactories::SelectActive(options_, write_schema_, pool_)); + PAIMON_ASSIGN_OR_RAISE(std::optional veto_reason, + GetDictionaryPassthroughVetoReason(plan_factory)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + CreateFilesReader(partition, bucket, dv_factory, to_compact, veto_reason)); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr data_file_path_factory, file_store_path_factory_->CreateDataFilePathFactory(partition, bucket)); - PAIMON_ASSIGN_OR_RAISE( - WriterFactory writer_factory, - GetDataFileWriterFactory(data_file_path_factory, write_schema_, write_cols_, to_compact)); + PAIMON_ASSIGN_OR_RAISE(WriterFactory writer_factory, + GetDataFileWriterFactory(data_file_path_factory, write_schema_, + write_cols_, to_compact, plan_factory)); + std::shared_ptr logical_type = arrow::struct_(write_schema_->fields()); + // Buffers allocated through the adaptor keep a raw pointer to it, and the writer may still + // hold a decoded column in its buffered row group, so it has to outlive the whole rewrite. + std::unique_ptr arrow_pool = GetArrowPool(pool_); auto rewriter = std::make_unique>>( options_.GetTargetFileSize(/*has_primary_key=*/false), @@ -188,6 +218,17 @@ Result>> AppendOnlyFileStoreWrite::Com auto struct_array = checked_pointer_cast(arrow_array); PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::RemoveFieldFromStructArray( struct_array, SpecialFields::ValueKind().Name())); + // The export below drops the type, leaving the writer to recover each column's encoding + // from the layout alone. Decode here, while the type is still known, every dictionary a + // veto forbids this rewrite from preserving; otherwise decode only shapes the layout + // cannot describe, such as the `dictionary(int64, large_utf8)` an ORC reader hands over + // under lazy decoding. A veto turns off the Parquet option, but another format's reader + // may emit dictionaries independently of it. + PAIMON_ASSIGN_OR_RAISE( + struct_array, + ArrowUtils::FlattenUnresolvableDictionaries( + struct_array, logical_type, arrow_pool.get(), + /*preserve_layout_recoverable_dictionaries=*/!veto_reason.has_value())); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*struct_array, c_array.get(), c_schema.get())); ArrowSchemaRelease(c_schema.get()); @@ -264,10 +305,9 @@ Result AppendOnlyFileStoreWrite::GetDat const std::shared_ptr& data_file_path_factory, const std::shared_ptr& schema, const std::optional>& write_cols, - const std::vector>& to_compact) const { + const std::vector>& to_compact, + const std::shared_ptr& plan_factory) const { auto seq_num_counter = std::make_shared(to_compact[0]->min_sequence_number); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr plan_factory, - ShreddingWritePlanFactories::SelectActive(options_, schema, pool_)); if (plan_factory != nullptr) { return std::make_shared( options_, table_schema_->Id(), schema, write_cols, seq_num_counter, @@ -278,9 +318,28 @@ Result AppendOnlyFileStoreWrite::GetDat data_file_path_factory, pool_); } +Result> AppendOnlyFileStoreWrite::GetDictionaryPassthroughVetoReason( + const std::shared_ptr& plan_factory) const { + std::shared_ptr file_format = options_.GetFileFormat(); + if (!file_format || file_format->Identifier() != kParquetFormat) { + return std::optional("the table does not write Parquet"); + } + PAIMON_ASSIGN_OR_RAISE(bool enable_dictionary, OptionsUtils::GetValueFromMap( + options_.ToMap(), kParquetEnableDictionary, + kDefaultParquetEnableDictionary)); + if (!enable_dictionary) { + return std::optional("parquet.enable-dictionary is false"); + } + if (plan_factory != nullptr) { + return std::optional("the table is written through a shredding writer"); + } + return std::optional(); +} + Result> AppendOnlyFileStoreWrite::CreateFilesReader( const BinaryRow& partition, int32_t bucket, DeletionVector::Factory dv_factory, - const std::vector>& files) const { + const std::vector>& files, + const std::optional& veto_reason) const { ReadContextBuilder context_builder(root_path_); context_builder.SetOptions(options_.ToMap()) .WithFileSystem(options_.GetFileSystem()) @@ -290,6 +349,23 @@ Result> AppendOnlyFileStoreWrite::CreateFilesReader .WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, context_builder.Finish()); std::map options = options_.ToMap(); + // Only the veto is applied here: forwarding is the table's decision, taken through + // `parquet.read.enable-dictionary-passthrough`, but a writer that cannot take a + // dictionary-encoded batch must not receive one because the table happens to set that option. + if (veto_reason.has_value()) { + // Overriding a table option is worth a line, but only for a table that set it: + // otherwise every rewrite would report a decision nobody made. A malformed value only + // costs the log line - the override below replaces it either way, so there is nothing to + // fail the rewrite over. + Result requested = + OptionsUtils::GetValueFromMap(options, kParquetReadEnableDictionaryPassthrough, + kDefaultParquetReadEnableDictionaryPassthrough); + if (requested.ok() && requested.value()) { + PAIMON_LOG_DEBUG(logger_, "Ignoring %s for this compaction rewrite: %s", + kParquetReadEnableDictionaryPassthrough, veto_reason->c_str()); + } + options[kParquetReadEnableDictionaryPassthrough] = "false"; + } PAIMON_ASSIGN_OR_RAISE(std::shared_ptr internal_read_context, InternalReadContext::Create(read_context, table_schema_, options)); auto read = std::make_unique(file_store_path_factory_, internal_read_context, diff --git a/src/paimon/core/operation/append_only_file_store_write.h b/src/paimon/core/operation/append_only_file_store_write.h index c6e5bf0b4..132118a72 100644 --- a/src/paimon/core/operation/append_only_file_store_write.h +++ b/src/paimon/core/operation/append_only_file_store_write.h @@ -66,6 +66,7 @@ class Executor; class Logger; class MemoryPool; class SchemaManager; +class ShreddingWritePlanFactory; class TableSchema; class IOManager; @@ -118,15 +119,41 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { return realtime_context_ != nullptr; } + /// @param plan_factory The active shredding write plan, or nullptr when the rewrite stays a + /// plain passthrough. Resolved by the caller because + /// `GetDictionaryPassthroughVetoReason` needs the same answer. Result GetDataFileWriterFactory( const std::shared_ptr& data_file_path_factory, const std::shared_ptr& schema, const std::optional>& write_cols, - const std::vector>& to_compact) const; + const std::vector>& to_compact, + const std::shared_ptr& plan_factory) const; + /// @param veto_reason What GetDictionaryPassthroughVetoReason() returned. Result> CreateFilesReader( const BinaryRow& partition, int32_t bucket, DeletionVector::Factory dv_factory, - const std::vector>& files) const; + const std::vector>& files, + const std::optional& veto_reason) const; + + /// Why `CompactRewrite` must not forward the dictionary encoding of its input files, or + /// `std::nullopt` when the table's `parquet.read.enable-dictionary-passthrough` may stand. + /// + /// A veto, not a decision: `std::nullopt` enables nothing. The option is off by default and + /// nothing here turns it on, because forwarding trades compaction CPU for output size and + /// which way that goes depends on the data. The "Dictionary Passthrough" section of + /// `docs/source/user_guide/compaction.rst` is where that trade is spelled out for users. + /// + /// The veto stands unless all three of: + /// + /// - a Parquet output file, since no other writer takes a dictionary-encoded batch; + /// - `parquet.enable-dictionary`, or the writer densifies what the reader just handed it and + /// the encoding is carried across the rewrite for nothing; + /// - a rewrite that stays a passthrough, since a shredding writer reshapes each batch against + /// a fixed physical schema and cannot take a dictionary-encoded one. + /// + /// @return The reason, meant for a log line rather than for branching on. + Result> GetDictionaryPassthroughVetoReason( + const std::shared_ptr& plan_factory) const; std::optional> write_cols_; std::shared_ptr realtime_context_; diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index 7605c4242..680c0907c 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -21,7 +21,10 @@ #include #include #include +#include +#include #include +#include #include "arrow/acero/options.h" #include "arrow/array/array_nested.h" @@ -40,6 +43,7 @@ #include "arrow/util/thread_pool.h" #include "fmt/format.h" #include "paimon/common/metrics/metrics_impl.h" +#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/status_utils.h" #include "paimon/common/utils/checked_cast.h" #include "paimon/common/utils/options_utils.h" @@ -55,7 +59,10 @@ #include "paimon/reader/batch_reader.h" #include "paimon/utils/roaring_bitmap32.h" #include "parquet/arrow/reader.h" +#include "parquet/metadata.h" #include "parquet/properties.h" +#include "parquet/schema.h" +#include "parquet/types.h" namespace arrow { class MemoryPool; @@ -132,6 +139,37 @@ bool HasSameNestedProjectionShape(const std::shared_ptr& read_t } } +// Whether every data page of `chunk` is dictionary-encoded, which is what makes reading the +// column as an Arrow DictionaryArray free. Writers that predate encoding statistics report none, +// and a chunk that fell back to PLAIN mid-way still carries the dictionary page it had already +// emitted, so an absent or mixed report means "not a passthrough candidate". +bool IsChunkFullyDictionaryEncoded(const ::parquet::ColumnChunkMetaData& chunk) { + if (!chunk.has_dictionary_page()) { + return false; + } + const std::vector<::parquet::PageEncodingStats>& encoding_stats = chunk.encoding_stats(); + if (encoding_stats.empty()) { + return false; + } + bool has_data_page = false; + for (const ::parquet::PageEncodingStats& stats : encoding_stats) { + if (stats.page_type != ::parquet::PageType::DATA_PAGE && + stats.page_type != ::parquet::PageType::DATA_PAGE_V2) { + continue; + } + if (stats.count <= 0) { + continue; + } + has_data_page = true; + // PLAIN_DICTIONARY is how a v1 writer spells dictionary indices; RLE_DICTIONARY is v2. + if (stats.encoding != ::parquet::Encoding::RLE_DICTIONARY && + stats.encoding != ::parquet::Encoding::PLAIN_DICTIONARY) { + return false; + } + } + return has_data_page; +} + // Resolve whether parquet-level pre-buffering should be enabled. When the framework // provides runtime hints, they describe the authoritative state of this read: once the // shared read-ahead cache takes over prefetching, disable parquet's own pre-buffering so @@ -149,16 +187,107 @@ ParquetFileBatchReader::ParquetFileBatchReader( std::shared_ptr&& input_stream, std::unique_ptr&& reader, const std::map& options, const std::shared_ptr& arrow_pool, - std::shared_ptr> storage_read_bytes) + std::shared_ptr> storage_read_bytes, + std::set dictionary_fields) : options_(options), arrow_pool_(arrow_pool), input_stream_(std::move(input_stream)), reader_(std::move(reader)), + dictionary_fields_(std::move(dictionary_fields)), read_ranges_(reader_->GetAllRowGroupRanges()), metrics_(std::make_shared()), storage_read_bytes_(std::move(storage_read_bytes)), logger_(Logger::GetLogger("ParquetFileBatchReader")) {} +std::set ParquetFileBatchReader::ResolveFullyDictionaryEncodedColumns( + const ::parquet::FileMetaData& metadata) { + std::set columns; + if (metadata.num_row_groups() == 0) { + return columns; + } + const ::parquet::SchemaDescriptor* schema = metadata.schema(); + for (int32_t i = 0; i < schema->num_columns(); ++i) { + // Arrow only reads BYTE_ARRAY leaves as dictionaries, and only a top-level column can be + // forwarded to the writer without rebuilding the nesting around it. + // + // `is_string()` narrows that further to STRING, leaving out the other BYTE_ARRAY leaf, + // BINARY. This is the reader's restriction, not the format's: the writer takes + // `dictionary(int32, binary)` and ArrowUtils::IsDictionaryLayoutRecoverableValueType() + // accepts it, but the option applies to every read of the table and the value accessors + // cannot read one. ColumnarUtils::GetView() asserts on a dictionary whose values are + // neither STRING nor LARGE_STRING and returns an empty view in a release build, and + // LiteralConverter rejects it. Every consumer here understands a STRING dictionary because + // the ORC reader has always produced one under lazy decoding; none was ever handed a + // BINARY one. Widening this needs those consumers first, not just the gate. + if (schema->Column(i)->physical_type() == ::parquet::Type::BYTE_ARRAY && + schema->Column(i)->logical_type()->is_string() && + schema->GetColumnRoot(i)->is_primitive()) { + columns.insert(i); + } + } + // Drop every candidate whose chunks are not dictionary-encoded end to end. A dictionary page + // alone does not say that: once the dictionary outgrows its page limit the writer emits the + // page it has and falls back to PLAIN for the rest, so a high-cardinality column keeps a + // dictionary page it no longer uses. Reading that as a dictionary would hash the PLAIN values + // back into a large in-memory dictionary, which is the work passthrough exists to avoid. + // Row groups are the outer loop so their metadata is materialized once each. + for (int32_t row_group = 0; row_group < metadata.num_row_groups() && !columns.empty(); + ++row_group) { + std::unique_ptr<::parquet::RowGroupMetaData> row_group_metadata = + metadata.RowGroup(row_group); + for (auto it = columns.begin(); it != columns.end();) { + if (IsChunkFullyDictionaryEncoded(*row_group_metadata->ColumnChunk(*it))) { + ++it; + } else { + it = columns.erase(it); + } + } + } + return columns; +} + +Result> ParquetFileBatchReader::GetLogicalFileSchema() const { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_schema, reader_->GetSchema()); + if (dictionary_fields_.empty()) { + return file_schema; + } + arrow::FieldVector fields; + fields.reserve(file_schema->num_fields()); + for (const auto& field : file_schema->fields()) { + if (field->type()->id() == arrow::Type::DICTIONARY) { + const auto& dictionary_type = + checked_cast(*field->type()); + fields.push_back(field->WithType(dictionary_type.value_type())); + } else { + fields.push_back(field); + } + } + return arrow::schema(fields, file_schema->metadata()); +} + +std::shared_ptr ParquetFileBatchReader::ApplyDictionaryReadTypes( + const std::shared_ptr& read_schema) const { + if (dictionary_fields_.empty()) { + return arrow::struct_(read_schema->fields()); + } + arrow::FieldVector fields; + fields.reserve(read_schema->num_fields()); + for (const auto& field : read_schema->fields()) { + // A passthrough column is always read at its file type, since this reader only ever casts + // timestamps, so checking the read type is the same as checking the file type. The + // predicate is the one the writer applies on the other side of the C data interface: both + // ends have to agree on which encodings survive the round trip, or a column this hands on + // encoded is one the writer refuses. + if (dictionary_fields_.count(field->name()) > 0 && + ArrowUtils::IsDictionaryLayoutRecoverableValueType(*field->type())) { + fields.push_back(field->WithType(arrow::dictionary(arrow::int32(), field->type()))); + } else { + fields.push_back(field); + } + } + return arrow::struct_(fields); +} + Result> ParquetFileBatchReader::Create( std::shared_ptr&& input_stream, const std::map& options, int32_t batch_size, @@ -177,6 +306,17 @@ Result> ParquetFileBatchReader::Create( PAIMON_RETURN_NOT_OK_FROM_ARROW( file_reader_builder.Open(input_stream, reader_properties, std::move(file_metadata))); + PAIMON_ASSIGN_OR_RAISE(bool enable_dictionary_passthrough, + OptionsUtils::GetValueFromMap( + options, PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH, + DEFAULT_PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH)); + if (enable_dictionary_passthrough) { + for (int32_t column_index : ResolveFullyDictionaryEncodedColumns( + *file_reader_builder.raw_reader()->metadata())) { + arrow_reader_properties.set_read_dictionary(column_index, /*read_dict=*/true); + } + } + std::unique_ptr<::parquet::arrow::FileReader> file_reader; PAIMON_RETURN_NOT_OK_FROM_ARROW(file_reader_builder.memory_pool(pool.get()) ->properties(arrow_reader_properties) @@ -184,9 +324,22 @@ Result> ParquetFileBatchReader::Create( PAIMON_ASSIGN_OR_RAISE(std::unique_ptr reader, FileReaderWrapper::Create(std::move(file_reader), static_cast(batch_size), pool)); - auto parquet_file_batch_reader = std::unique_ptr( - new ParquetFileBatchReader(std::move(input_stream), std::move(reader), options, pool, - std::move(storage_read_bytes))); + // Arrow silently ignores set_read_dictionary for leaves it cannot read as dictionaries, + // so take the columns it really emits that way from the schema it just derived. + std::set dictionary_fields; + if (enable_dictionary_passthrough) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr reader_schema, + reader->GetSchema()); + for (const auto& field : reader_schema->fields()) { + if (field->type()->id() == arrow::Type::DICTIONARY) { + dictionary_fields.insert(field->name()); + } + } + } + auto parquet_file_batch_reader = + std::unique_ptr(new ParquetFileBatchReader( + std::move(input_stream), std::move(reader), options, pool, + std::move(storage_read_bytes), std::move(dictionary_fields))); PAIMON_ASSIGN_OR_RAISE(std::unique_ptr<::ArrowSchema> file_schema, parquet_file_batch_reader->GetFileSchema()); PAIMON_RETURN_NOT_OK(parquet_file_batch_reader->SetReadSchema( @@ -198,7 +351,7 @@ Result> ParquetFileBatchReader::Create( Result> ParquetFileBatchReader::GetFileSchema() const { try { - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_schema, reader_->GetSchema()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_schema, GetLogicalFileSchema()); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr new_schema, ParquetFieldIdConverter::GetPaimonIdsFromParquetIds(file_schema)); PAIMON_ASSIGN_OR_RAISE( @@ -222,7 +375,7 @@ Status ParquetFileBatchReader::SetReadSchema( PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr read_schema, arrow::ImportSchema(schema)); - PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_schema, reader_->GetSchema()); + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr file_schema, GetLogicalFileSchema()); // Recursively match read_schema against file_schema by field names. // STRUCT supports sub-field projection; LIST/MAP require exact type match. @@ -301,7 +454,7 @@ Status ParquetFileBatchReader::SetReadSchema( } } - read_data_type_ = arrow::struct_(read_schema->fields()); + read_data_type_ = ApplyDictionaryReadTypes(read_schema); metrics_->SetCounter(ParquetMetrics::READ_ROW_GROUPS_TOTAL, reader_->GetNumberOfRowGroups()); diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.h b/src/paimon/format/parquet/parquet_file_batch_reader.h index 2b1097fb4..7c8e3bbdf 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.h +++ b/src/paimon/format/parquet/parquet_file_batch_reader.h @@ -162,13 +162,31 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { std::unique_ptr&& reader, const std::map& options, const std::shared_ptr& arrow_pool, - std::shared_ptr> storage_read_bytes); + std::shared_ptr> storage_read_bytes, + std::set dictionary_fields); static Result<::parquet::ArrowReaderProperties> CreateArrowReaderProperties( const std::shared_ptr& pool, const std::map& options, int32_t batch_size, const std::optional& hints); + /// Leaf column indices that are candidates for `set_read_dictionary`: non-nested STRING + /// columns whose every data page, in every row group, is dictionary-encoded. The definition + /// says why STRING and not every BYTE_ARRAY leaf. + static std::set ResolveFullyDictionaryEncodedColumns( + const ::parquet::FileMetaData& metadata); + + /// The file schema with the dictionary encoding of the passthrough columns removed. That + /// encoding describes the batches this reader emits, not the types stored in the file, so + /// everything reasoning about the file's types (projection, predicate binding) sees the + /// logical schema and only `read_data_type_` carries the dictionaries. + Result> GetLogicalFileSchema() const; + + /// Builds the read type from `read_schema`, re-applying the dictionary encoding of every + /// passthrough column so the read type keeps describing what `NextBatch()` produces. + std::shared_ptr ApplyDictionaryReadTypes( + const std::shared_ptr& read_schema) const; + static void FlattenSchema(const std::shared_ptr& type, int32_t* index, std::vector* index_vector) { if (type->id() == arrow::Type::STRUCT || type->id() == arrow::Type::LIST || @@ -262,6 +280,10 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { std::shared_ptr read_data_type_; + // Top-level file fields emitted as Arrow DictionaryArray. Empty unless + // PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH is on. + std::set dictionary_fields_; + std::vector> read_ranges_; std::shared_ptr metrics_; diff --git a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp index 1409f24bf..e64fe3650 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -35,10 +35,13 @@ #include "arrow/array/builder_primitive.h" #include "arrow/c/abi.h" #include "arrow/c/bridge.h" +#include "arrow/compute/api.h" #include "arrow/io/caching.h" +#include "arrow/io/file.h" #include "arrow/io/interfaces.h" #include "arrow/ipc/api.h" #include "arrow/ipc/json_simple.h" +#include "fmt/format.h" #include "gtest/gtest.h" #include "paimon/common/io/cache_input_stream.h" #include "paimon/common/metrics/metrics_impl.h" @@ -67,7 +70,9 @@ #include "paimon/testing/utils/testharness.h" #include "paimon/testing/utils/timezone_guard.h" #include "paimon/utils/roaring_bitmap32.h" +#include "parquet/arrow/reader.h" #include "parquet/file_reader.h" +#include "parquet/metadata.h" #include "parquet/properties.h" namespace paimon { @@ -1730,4 +1735,302 @@ TEST_F(ParquetFileBatchReaderTest, TestPreBufferRangeFeedsReadAheadCache) { ASSERT_EQ(miss_bytes, baseline_miss_bytes); } +struct DictionaryPassthroughResult { + // First batch of a {f4: int32, f8: utf8} projection. + std::shared_ptr batch; + // Struct type the reader reports for the whole file, all 13 fields in file order. + std::shared_ptr file_type; +}; + +TEST_F(ParquetFileBatchReaderTest, TestDictionaryPassthrough) { + // `std::nullopt` leaves the option out of the map entirely, which is what an ordinary read + // does and what has to keep resolving to "off". + auto read_projection = [&](std::optional enable_dictionary_passthrough, + bool enable_dictionary_on_write) { + WriteArray(file_path_, struct_array_, schema_, + /*write_batch_size=*/struct_array_->length(), enable_dictionary_on_write, + /*max_row_group_length=*/struct_array_->length()); + + EXPECT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs_->Open(file_path_)); + auto length = fs_->GetFileStatus(file_path_).value().GetLen(); + auto in_stream = + std::make_unique(std::move(input_stream), length, pool_); + std::map options; + if (enable_dictionary_passthrough.has_value()) { + options[PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH] = + *enable_dictionary_passthrough ? "true" : "false"; + } + auto read_schema = + MakeReadSchema({arrow::field("f4", arrow::int32()), arrow::field("f8", arrow::utf8())}); + auto reader = PrepareParquetFileBatchReader(std::move(in_stream), options, read_schema, + /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, batch_size_); + + DictionaryPassthroughResult result; + EXPECT_OK_AND_ASSIGN(std::unique_ptr c_file_schema, reader->GetFileSchema()); + result.file_type = arrow::ImportType(c_file_schema.get()).ValueOrDie(); + + EXPECT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + EXPECT_FALSE(BatchReader::IsEofBatch(batch)); + auto& [c_array, c_schema] = batch; + result.batch = checked_pointer_cast( + arrow::ImportArray(c_array.get(), c_schema.get()).ValueOrDie()); + reader->Close(); + return result; + }; + + { + // The column is handed over encoded while the file schema keeps reporting the logical + // type: the dictionary describes the emitted batches, not the file. + DictionaryPassthroughResult result = read_projection(/*enable_dictionary_passthrough=*/true, + /*enable_dictionary_on_write=*/true); + ASSERT_EQ(arrow::Type::STRING, result.file_type->field(7)->type()->id()); + ASSERT_EQ(arrow::Type::INT32, result.batch->field(0)->type()->id()); + ASSERT_EQ(arrow::Type::DICTIONARY, result.batch->field(1)->type()->id()); + auto decoded = + arrow::compute::Cast(result.batch->field(1), arrow::utf8()).ValueOrDie().make_array(); + auto strings = checked_pointer_cast(decoded); + ASSERT_EQ(struct_array_->length(), strings->length()); + for (int64_t i = 0; i < strings->length(); ++i) { + ASSERT_EQ(fmt::format("s3{}", i + 1), strings->GetString(i)); + } + } + { + // The file has no dictionary pages, so the gate keeps the column materialized instead of + // moving the hashing from the writer to the reader. + DictionaryPassthroughResult result = read_projection(/*enable_dictionary_passthrough=*/true, + /*enable_dictionary_on_write=*/false); + ASSERT_EQ(arrow::Type::STRING, result.batch->field(1)->type()->id()); + } + { + // Explicitly off: the kill switch a table can set to opt the rewrite out. + DictionaryPassthroughResult result = + read_projection(/*enable_dictionary_passthrough=*/false, + /*enable_dictionary_on_write=*/true); + ASSERT_EQ(arrow::Type::STRING, result.batch->field(1)->type()->id()); + } + { + // Absent, which is how this reader is reached on a table that never set the option - + // most of them. The default has to be off, or an ordinary scan would start emitting + // dictionary batches at consumers that do not unwrap them. A table that does set it gets + // them on every read, not only in the compaction rewrite. + DictionaryPassthroughResult result = + read_projection(/*enable_dictionary_passthrough=*/std::nullopt, + /*enable_dictionary_on_write=*/true); + ASSERT_EQ(arrow::Type::STRING, result.batch->field(1)->type()->id()); + } +} + +TEST_F(ParquetFileBatchReaderTest, TestDictionaryPassthroughSkipsBinaryColumn) { + // Parquet stores STRING and BINARY in the same BYTE_ARRAY leaf and dictionary-encodes both, so + // the gate has to exclude BINARY by logical type. It does, because nothing downstream can read + // `dictionary(int32, binary)`: ColumnarUtils::GetView() asserts on it and returns an empty view + // in a release build, and LiteralConverter rejects it. `f8` is the control - same physical + // type, same pages, and it is forwarded - so this fails if the exclusion is ever widened back. + WriteArray(file_path_, struct_array_, schema_, + /*write_batch_size=*/struct_array_->length(), /*enable_dictionary=*/true, + /*max_row_group_length=*/struct_array_->length()); + + ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs_->Open(file_path_)); + auto length = fs_->GetFileStatus(file_path_).value().GetLen(); + auto in_stream = + std::make_unique(std::move(input_stream), length, pool_); + std::map options; + options[PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH] = "true"; + auto read_schema = + MakeReadSchema({arrow::field("f8", arrow::utf8()), arrow::field("f9", arrow::binary())}); + auto reader = PrepareParquetFileBatchReader(std::move(in_stream), options, read_schema, + /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, batch_size_); + + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + auto& [c_array, c_schema] = batch; + auto read_array = checked_pointer_cast( + arrow::ImportArray(c_array.get(), c_schema.get()).ValueOrDie()); + ASSERT_EQ(arrow::Type::DICTIONARY, read_array->field(0)->type()->id()) + << read_array->field(0)->type()->ToString(); + ASSERT_EQ(arrow::Type::BINARY, read_array->field(1)->type()->id()) + << read_array->field(1)->type()->ToString(); + // Materialized, and still the bytes the file holds. + ASSERT_EQ("a31", checked_pointer_cast(read_array->field(1))->GetString(0)); + reader->Close(); +} + +TEST_F(ParquetFileBatchReaderTest, TestDictionaryPassthroughSkipsFallbackToPlain) { + // A column whose dictionary outgrows its page limit keeps the dictionary page the writer had + // already emitted and encodes the rest as PLAIN. Reading that back as a dictionary would hash + // the plain values into a large in-memory dictionary, which is the work passthrough exists to + // avoid, so the gate has to look at the data page encodings rather than the dictionary page. + constexpr int32_t kRows = 4000; + arrow::StringBuilder value_builder; + for (int32_t i = 0; i < kRows; ++i) { + ASSERT_TRUE(value_builder.Append(fmt::format("unique_value_{}", i)).ok()); + } + std::shared_ptr values; + ASSERT_TRUE(value_builder.Finish(&values).ok()); + auto field = arrow::field("f0", arrow::utf8()); + auto write_schema = arrow::schema({field}); + auto struct_array = arrow::StructArray::Make({values}, {field}).ValueOrDie(); + + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_fallback.parquet"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs_->Create(file_path, /*overwrite=*/true)); + ::parquet::WriterProperties::Builder writer_builder; + writer_builder.enable_dictionary(); + // Small enough that the dictionary overflows partway through and the writer falls back. + writer_builder.dictionary_pagesize_limit(1024); + ASSERT_OK_AND_ASSIGN(auto format_writer, + ParquetFormatWriter::Create(out, write_schema, writer_builder.build(), + DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, pool_)); + auto arrow_array = std::make_unique(); + ASSERT_TRUE(arrow::ExportArray(*struct_array, arrow_array.get()).ok()); + ASSERT_OK(format_writer->AddBatch(arrow_array.get())); + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + + // The dictionary page written before the fallback is still in the chunk, which is exactly why + // its presence cannot be the signal; the data pages are what say the column went plain. Both + // are asserted so the test fails loudly if the fixture stops producing a mixed chunk rather + // than quietly passing for the wrong reason. + auto metadata_file = arrow::io::ReadableFile::Open(file_path, pool_.get()); + ASSERT_TRUE(metadata_file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> metadata_reader; + ASSERT_TRUE( + ::parquet::arrow::OpenFile(metadata_file.ValueOrDie(), pool_.get(), &metadata_reader).ok()); + std::unique_ptr<::parquet::ColumnChunkMetaData> column_chunk = + metadata_reader->parquet_reader()->metadata()->RowGroup(0)->ColumnChunk(0); + ASSERT_TRUE(column_chunk->has_dictionary_page()); + int32_t plain_data_pages = 0; + int32_t data_pages = 0; + for (const ::parquet::PageEncodingStats& stats : column_chunk->encoding_stats()) { + if (stats.page_type != ::parquet::PageType::DATA_PAGE && + stats.page_type != ::parquet::PageType::DATA_PAGE_V2) { + continue; + } + data_pages += stats.count; + if (stats.encoding == ::parquet::Encoding::PLAIN) { + plain_data_pages += stats.count; + } + } + ASSERT_GT(data_pages, 0); + ASSERT_GT(plain_data_pages, 0) << "fixture no longer falls back to plain encoding"; + + ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs_->Open(file_path)); + auto length = fs_->GetFileStatus(file_path).value().GetLen(); + auto in_stream = + std::make_unique(std::move(input_stream), length, pool_); + std::map options; + options[PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH] = "true"; + auto reader = PrepareParquetFileBatchReader(std::move(in_stream), options, write_schema, + /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, kRows); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + auto& [c_array, c_schema] = batch; + auto read_array = checked_pointer_cast( + arrow::ImportArray(c_array.get(), c_schema.get()).ValueOrDie()); + ASSERT_EQ(arrow::Type::STRING, read_array->field(0)->type()->id()); + reader->Close(); +} + +TEST_F(ParquetFileBatchReaderTest, TestDictionaryPassthroughRequiresEveryRowGroup) { + // The gate is a per-file decision taken over every row group, which is what the documentation + // promises. A file whose first row group is fully dictionary-encoded and whose second falls + // back to plain is the case that separates "checked the first row group" from "checked all of + // them": looking only at row group 0 would forward the column and then hash the plain values + // of row group 1 back into a large in-memory dictionary. + constexpr int32_t kRowsPerGroup = 64; + constexpr int32_t kWriteBatchSize = 16; + auto field = arrow::field("f0", arrow::utf8()); + auto write_schema = arrow::schema({field}); + + // Two distinct short values, so this row group stays dictionary-encoded end to end. + arrow::StringBuilder low_cardinality_builder; + for (int32_t i = 0; i < kRowsPerGroup; ++i) { + ASSERT_TRUE(low_cardinality_builder.Append(i % 2 == 0 ? "a" : "b").ok()); + } + // Distinct long values, so the dictionary outgrows its page limit after the first write batch + // and the rest of this row group is plain. + arrow::StringBuilder high_cardinality_builder; + for (int32_t i = 0; i < kRowsPerGroup; ++i) { + ASSERT_TRUE( + high_cardinality_builder.Append(fmt::format("distinct_value_padded_out_{}", i)).ok()); + } + std::shared_ptr low_cardinality, high_cardinality; + ASSERT_TRUE(low_cardinality_builder.Finish(&low_cardinality).ok()); + ASSERT_TRUE(high_cardinality_builder.Finish(&high_cardinality).ok()); + + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_mixed_row_groups.parquet"); + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs_->Create(file_path, /*overwrite=*/true)); + ::parquet::WriterProperties::Builder writer_builder; + writer_builder.enable_dictionary(); + writer_builder.dictionary_pagesize_limit(64); + // The dictionary limit is only checked once per write batch, so the fallback can only land + // mid-chunk if a row group spans several of them. + writer_builder.write_batch_size(kWriteBatchSize); + // One row group per AddBatch, which is what puts the two encodings in separate chunks. + writer_builder.max_row_group_length(kRowsPerGroup); + ASSERT_OK_AND_ASSIGN(auto format_writer, + ParquetFormatWriter::Create(out, write_schema, writer_builder.build(), + DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, pool_)); + for (const std::shared_ptr& values : {low_cardinality, high_cardinality}) { + auto struct_array = arrow::StructArray::Make({values}, {field}).ValueOrDie(); + auto arrow_array = std::make_unique(); + ASSERT_TRUE(arrow::ExportArray(*struct_array, arrow_array.get()).ok()); + ASSERT_OK(format_writer->AddBatch(arrow_array.get())); + } + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + + // Pin the fixture: without this the read assertion below would also pass on a file whose + // first row group was never dictionary-encoded in the first place. + auto metadata_file = arrow::io::ReadableFile::Open(file_path, pool_.get()); + ASSERT_TRUE(metadata_file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> metadata_reader; + ASSERT_TRUE( + ::parquet::arrow::OpenFile(metadata_file.ValueOrDie(), pool_.get(), &metadata_reader).ok()); + std::shared_ptr<::parquet::FileMetaData> metadata = + metadata_reader->parquet_reader()->metadata(); + ASSERT_EQ(2, metadata->num_row_groups()); + auto count_plain_data_pages = [](const ::parquet::ColumnChunkMetaData& chunk) { + int32_t plain = 0; + for (const ::parquet::PageEncodingStats& stats : chunk.encoding_stats()) { + if ((stats.page_type == ::parquet::PageType::DATA_PAGE || + stats.page_type == ::parquet::PageType::DATA_PAGE_V2) && + stats.encoding == ::parquet::Encoding::PLAIN) { + plain += stats.count; + } + } + return plain; + }; + ASSERT_EQ(0, count_plain_data_pages(*metadata->RowGroup(0)->ColumnChunk(0))) + << "first row group was expected to stay dictionary-encoded"; + ASSERT_GT(count_plain_data_pages(*metadata->RowGroup(1)->ColumnChunk(0)), 0) + << "second row group was expected to fall back to plain"; + + ASSERT_OK_AND_ASSIGN(std::shared_ptr input_stream, fs_->Open(file_path)); + auto length = fs_->GetFileStatus(file_path).value().GetLen(); + auto in_stream = + std::make_unique(std::move(input_stream), length, pool_); + std::map options; + options[PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH] = "true"; + auto reader = PrepareParquetFileBatchReader(std::move(in_stream), options, write_schema, + /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt, kRowsPerGroup); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + auto& [c_array, c_schema] = batch; + auto read_array = checked_pointer_cast( + arrow::ImportArray(c_array.get(), c_schema.get()).ValueOrDie()); + // One row group disqualifies the whole column, including the row group that did qualify. + ASSERT_EQ(arrow::Type::STRING, read_array->field(0)->type()->id()); + reader->Close(); +} + } // namespace paimon::parquet::test diff --git a/src/paimon/format/parquet/parquet_format_defs.h b/src/paimon/format/parquet/parquet_format_defs.h index 8b205a09a..d22506c45 100644 --- a/src/paimon/format/parquet/parquet_format_defs.h +++ b/src/paimon/format/parquet/parquet_format_defs.h @@ -102,12 +102,25 @@ static inline const char PARQUET_READ_ENABLE_PAGE_INDEX_FILTER[] = // Default is true. static inline const char PARQUET_READ_ENABLE_PRE_BUFFER[] = "parquet.read.enable-pre-buffer"; +// Emit dictionary-encoded STRING columns as Arrow DictionaryArray instead of one copy of the +// value per row. Which columns qualify, and why BINARY does not, is decided by +// ParquetFileBatchReader::ResolveFullyDictionaryEncodedColumns(). +// +// Off by default, and set on the table by a user who has measured the trade-off - see the +// "Dictionary Passthrough" section of `docs/source/user_guide/compaction.rst`. The append +// compaction rewrite is the only consumer that gains from it, and it can only veto the option, +// never turn it on. Being a read option it applies to every read of the table, so a consumer that +// reads values through its own accessor has to unwrap DictionaryArray. +static inline const char PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH[] = + "parquet.read.enable-dictionary-passthrough"; + static constexpr uint32_t DEFAULT_PARQUET_READ_CACHE_OPTION_PREFETCH_LIMIT = 0; // Default value of hole size limit, inherited from Arrow static constexpr uint32_t DEFAULT_PARQUET_READ_CACHE_OPTION_HOLE_SIZE_LIMIT = 8 * 1024; static constexpr uint32_t DEFAULT_PARQUET_READ_CACHE_OPTION_RANGE_SIZE_LIMIT = 32 * 1024 * 1024; static constexpr uint32_t DEFAULT_PARQUET_READ_PREDICATE_NODE_COUNT_LIMIT = 512; static constexpr bool DEFAULT_PARQUET_READ_ENABLE_PAGE_INDEX_FILTER = true; +static constexpr bool DEFAULT_PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH = false; static constexpr char DEFAULT_PARQUET_READ_BITMAP_STRATEGY[] = "coalesce"; static constexpr uint32_t DEFAULT_PARQUET_READ_ROW_RANGES_COALESCE_HOLE_SIZE_LIMIT = 32; diff --git a/src/paimon/format/parquet/parquet_format_writer.cpp b/src/paimon/format/parquet/parquet_format_writer.cpp index 0a8e38b43..768615a97 100644 --- a/src/paimon/format/parquet/parquet_format_writer.cpp +++ b/src/paimon/format/parquet/parquet_format_writer.cpp @@ -23,14 +23,20 @@ #include #include +#include "arrow/array/array_dict.h" #include "arrow/c/bridge.h" +#include "arrow/compute/cast.h" #include "arrow/memory_pool.h" #include "arrow/record_batch.h" +#include "arrow/type.h" #include "arrow/util/base64.h" #include "arrow/util/key_value_metadata.h" #include "paimon/common/metrics/metrics_impl.h" #include "paimon/common/utils/arrow/arrow_output_stream_adapter.h" +#include "paimon/common/utils/arrow/arrow_utils.h" #include "paimon/common/utils/arrow/status_utils.h" +#include "paimon/common/utils/checked_cast.h" +#include "paimon/core/casting/casting_utils.h" #include "paimon/format/parquet/parquet_format_defs.h" #include "parquet/arrow/writer.h" #include "parquet/properties.h" @@ -64,8 +70,10 @@ Result> ParquetFormatWriter::Create( } Status ParquetFormatWriter::AddBatch(ArrowArray* batch) { + PAIMON_ASSIGN_OR_RAISE(std::shared_ptr batch_schema, ResolveBatchSchema(batch)); PAIMON_ASSIGN_OR_RAISE_FROM_ARROW(std::shared_ptr<::arrow::RecordBatch> record_batch, - arrow::ImportRecordBatch(batch, schema_)); + arrow::ImportRecordBatch(batch, batch_schema)); + PAIMON_ASSIGN_OR_RAISE(record_batch, FlattenUnwritableDictionaries(record_batch)); if (static_cast(pool_->bytes_allocated()) > max_memory_use_) { PAIMON_RETURN_NOT_OK_FROM_ARROW(writer_->NewBufferedRowGroup()); } @@ -74,6 +82,46 @@ Status ParquetFormatWriter::AddBatch(ArrowArray* batch) { return Status::OK(); } +Result> ParquetFormatWriter::ResolveBatchSchema( + const ::ArrowArray* batch) const { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr batch_type, + ArrowUtils::ResolveDictionaryStructTypeFromLayout(logical_struct_type_, batch)); + if (batch_type == logical_struct_type_) { + return schema_; + } + return arrow::schema(batch_type->fields(), schema_->metadata()); +} + +Result> ParquetFormatWriter::FlattenUnwritableDictionaries( + const std::shared_ptr& record_batch) const { + arrow::ArrayVector columns; + arrow::FieldVector fields; + for (int32_t i = 0; i < record_batch->num_columns(); ++i) { + const std::shared_ptr& column = record_batch->column(i); + if (column->type_id() != arrow::Type::DICTIONARY || + checked_cast(*column).dictionary()->null_count() == 0) { + continue; + } + if (columns.empty()) { + columns = record_batch->columns(); + fields = record_batch->schema()->fields(); + } + const auto& dictionary_type = checked_cast(*column->type()); + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr flattened, + CastingUtils::Cast(column, dictionary_type.value_type(), + arrow::compute::CastOptions::Safe(), pool_.get())); + columns[i] = std::move(flattened); + fields[i] = fields[i]->WithType(dictionary_type.value_type()); + } + if (columns.empty()) { + return record_batch; + } + return arrow::RecordBatch::Make(arrow::schema(fields, record_batch->schema()->metadata()), + record_batch->num_rows(), std::move(columns)); +} + Status ParquetFormatWriter::Flush() { metrics_->SetCounter(ParquetMetrics::WRITE_RECORD_COUNT, total_records_written_); return Status::OK(); @@ -119,6 +167,7 @@ ParquetFormatWriter::ParquetFormatWriter(std::unique_ptr<::parquet::arrow::FileW out_(out), writer_(std::move(writer)), schema_(schema), + logical_struct_type_(arrow::struct_(schema->fields())), metrics_(std::make_shared()), max_memory_use_(max_memory_use) {} diff --git a/src/paimon/format/parquet/parquet_format_writer.h b/src/paimon/format/parquet/parquet_format_writer.h index 4ab58d73c..48f8253de 100644 --- a/src/paimon/format/parquet/parquet_format_writer.h +++ b/src/paimon/format/parquet/parquet_format_writer.h @@ -31,7 +31,9 @@ #include "parquet/arrow/writer.h" namespace arrow { +class DataType; class MemoryPool; +class RecordBatch; class Schema; } // namespace arrow namespace paimon { @@ -76,10 +78,30 @@ class ParquetFormatWriter : public FormatWriter { Result GetEstimateLength() const; + /// Returns the schema describing `batch` as it is laid out, which is `schema_` unless some of + /// its columns arrived dictionary-encoded. The Parquet write schema stays `schema_` either + /// way, since a dictionary is just an encoding of the same logical column. + /// + /// `parquet::arrow::FileWriter` writes the first dictionary a column presents in a row group + /// through `WriteArrowDictionary()`, without materializing its values. It keeps only that one: + /// a later batch carrying a different dictionary makes the column fall back to plain encoding + /// for the rest of the row group, so the values still round-trip but the output stops being + /// dictionary-encoded there. Passing an encoding on therefore saves work on the way in, not + /// necessarily on the way out. + Result> ResolveBatchSchema(const ::ArrowArray* batch) const; + + /// Flattens, per column, the dictionaries that Arrow's Parquet writer rejects outright, so + /// one such column does not fail the whole batch. Currently only dictionaries holding nulls + /// in their values, the one case `parquet::arrow` does not fall back on by itself. + Result> FlattenUnwritableDictionaries( + const std::shared_ptr& record_batch) const; + std::shared_ptr pool_; std::shared_ptr out_; std::unique_ptr<::parquet::arrow::FileWriter> writer_; std::shared_ptr schema_; + // Struct view of schema_, matched against the layout of each incoming batch. + std::shared_ptr logical_struct_type_; std::shared_ptr metrics_; int64_t total_records_written_ = 0; uint64_t max_memory_use_; diff --git a/src/paimon/format/parquet/parquet_format_writer_test.cpp b/src/paimon/format/parquet/parquet_format_writer_test.cpp index 117a1450b..3d1c00872 100644 --- a/src/paimon/format/parquet/parquet_format_writer_test.cpp +++ b/src/paimon/format/parquet/parquet_format_writer_test.cpp @@ -19,12 +19,14 @@ #include "paimon/format/parquet/parquet_format_writer.h" #include +#include #include #include #include #include "arrow/api.h" #include "arrow/array/array_binary.h" +#include "arrow/array/array_dict.h" #include "arrow/array/array_primitive.h" #include "arrow/array/builder_binary.h" #include "arrow/array/builder_nested.h" @@ -32,9 +34,11 @@ #include "arrow/c/abi.h" #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" +#include "arrow/compute/api.h" #include "arrow/io/file.h" #include "arrow/ipc/api.h" #include "arrow/memory_pool.h" +#include "fmt/format.h" #include "gtest/gtest.h" #include "paimon/common/utils/arrow/mem_utils.h" #include "paimon/common/utils/checked_cast.h" @@ -49,12 +53,14 @@ #include "paimon/memory/memory_pool.h" #include "paimon/metrics.h" #include "paimon/record_batch.h" +#include "paimon/status.h" #include "paimon/testing/utils/testharness.h" #include "parquet/arrow/reader.h" #include "parquet/file_reader.h" #include "parquet/metadata.h" #include "parquet/properties.h" #include "parquet/schema.h" +#include "parquet/statistics.h" namespace arrow { class Array; @@ -184,6 +190,131 @@ class ParquetFormatWriterTest : public ::testing::Test { } } + /// Builds the {col1, col2, col3} batch of this fixture with a low-cardinality col1 that is + /// either dictionary-encoded or flat, so the same rows reach the writer in both encodings. + /// Row i holds col1 = "", col2 = i and col3 = i % 2. + std::shared_ptr PrepareEncodedArray( + int32_t record_batch_size, int32_t offset, bool dictionary_encoded, + bool null_in_dictionary = false, const std::string& dictionary_prefix = "dict_") const { + arrow::StringBuilder dictionary_builder; + for (int32_t i = 0; i < 3; ++i) { + if (null_in_dictionary && i == 1) { + EXPECT_TRUE(dictionary_builder.AppendNull().ok()); + } else { + EXPECT_TRUE( + dictionary_builder.Append(fmt::format("{}{}", dictionary_prefix, i)).ok()); + } + } + std::shared_ptr dictionary; + EXPECT_TRUE(dictionary_builder.Finish(&dictionary).ok()); + + arrow::Int32Builder index_builder; + arrow::Int32Builder int_builder; + arrow::BooleanBuilder bool_builder; + for (int32_t i = offset; i < offset + record_batch_size; ++i) { + EXPECT_TRUE(index_builder.Append(i % 3).ok()); + EXPECT_TRUE(int_builder.Append(i).ok()); + EXPECT_TRUE(bool_builder.Append(static_cast(i % 2)).ok()); + } + std::shared_ptr indices, int_array, bool_array; + EXPECT_TRUE(index_builder.Finish(&indices).ok()); + EXPECT_TRUE(int_builder.Finish(&int_array).ok()); + EXPECT_TRUE(bool_builder.Finish(&bool_array).ok()); + + auto dictionary_array = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::utf8()), + indices, dictionary) + .ValueOrDie(); + std::shared_ptr string_array = dictionary_array; + if (!dictionary_encoded) { + string_array = + arrow::compute::Cast(dictionary_array, arrow::utf8()).ValueOrDie().make_array(); + } + return arrow::StructArray::Make({string_array, int_array, bool_array}, + std::vector{"col1", "col2", "col3"}) + .ValueOrDie(); + } + + void AddStructArrayOnce(const std::shared_ptr& format_writer, + const std::shared_ptr& array) const { + auto arrow_array = std::make_unique(); + ASSERT_TRUE(arrow::ExportArray(*array, arrow_array.get()).ok()); + ASSERT_OK(format_writer->AddBatch(arrow_array.get())); + } + + /// Fraction of `chunk`'s data pages written with a dictionary index encoding. A dictionary + /// page on its own says nothing - the writer emits one and then falls back to plain when the + /// dictionary it kept no longer matches - so the encoding of the data pages is what tells + /// whether the column really came out dictionary-encoded. + static std::pair CountDictionaryDataPages( + const ::parquet::ColumnChunkMetaData& chunk) { + int32_t dictionary_pages = 0; + int32_t data_pages = 0; + for (const ::parquet::PageEncodingStats& stats : chunk.encoding_stats()) { + if (stats.page_type != ::parquet::PageType::DATA_PAGE && + stats.page_type != ::parquet::PageType::DATA_PAGE_V2) { + continue; + } + data_pages += stats.count; + if (stats.encoding == ::parquet::Encoding::RLE_DICTIONARY || + stats.encoding == ::parquet::Encoding::PLAIN_DICTIONARY) { + dictionary_pages += stats.count; + } + } + return {dictionary_pages, data_pages}; + } + + void CheckEncodedResult(const std::string& file_path, int32_t row_count, + bool null_in_dictionary) const { + auto file = arrow::io::ReadableFile::Open(file_path, arrow_pool_.get()); + ASSERT_TRUE(file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> reader; + auto status = ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &reader); + ASSERT_TRUE(status.ok()) << status.ToString(); + const ::parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get(); + ASSERT_EQ(metadata->num_rows(), row_count); + // Whether the values arrived encoded or flat, every data page comes out dictionary-encoded. + auto [dictionary_pages, data_pages] = + CountDictionaryDataPages(*metadata->RowGroup(0)->ColumnChunk(0)); + ASSERT_GT(data_pages, 0); + ASSERT_EQ(data_pages, dictionary_pages); + + std::shared_ptr<::arrow::ChunkedArray> col0_array; + ASSERT_TRUE(reader->ReadColumn(0, &col0_array).ok()); + int32_t row = 0; + for (const auto& chunk : col0_array->chunks()) { + const auto& string_array = checked_pointer_cast(chunk); + ASSERT_TRUE(string_array); + for (int64_t i = 0; i < string_array->length(); ++i, ++row) { + if (null_in_dictionary && row % 3 == 1) { + ASSERT_TRUE(string_array->IsNull(i)); + } else { + ASSERT_EQ(fmt::format("dict_{}", row % 3), string_array->GetString(i)); + } + } + } + ASSERT_EQ(row_count, row); + } + + /// @param max_memory_use Lower it to make every AddBatch start a fresh buffered row group, + /// which flushes the previous one to the output stream. + std::shared_ptr CreateEncodedWriter( + const std::string& file_path, std::shared_ptr* out, + uint64_t max_memory_use = DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE) const { + EXPECT_OK_AND_ASSIGN(*out, fs_->Create(file_path, /*overwrite=*/false)); + ::parquet::WriterProperties::Builder builder; + builder.enable_dictionary(); + // What ParquetWriterBuilder does in production. Without it the encoders allocate from + // Arrow's default pool, so `max_memory_use` would be compared against a pool the writer + // never touches and the row-group rotation it drives would never fire. + builder.memory_pool(arrow_pool_.get()); + EXPECT_OK_AND_ASSIGN( + std::shared_ptr format_writer, + ParquetFormatWriter::Create(*out, PrepareArrowSchema().first, builder.build(), + max_memory_use, arrow_pool_)); + return format_writer; + } + private: std::unique_ptr dir_; std::shared_ptr fs_; @@ -467,4 +598,360 @@ TEST_F(ParquetFormatWriterTest, TestTimestampType) { ASSERT_OK(out->Close()); } +TEST_F(ParquetFormatWriterTest, TestWriteDictionaryEncodedColumn) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_passthrough"); + std::shared_ptr out; + std::shared_ptr format_writer = CreateEncodedWriter(file_path, &out); + + // The writer is created from the logical schema, which a dictionary-encoded batch no longer + // matches, and the encoding may alternate from one batch to the next. + AddStructArrayOnce(format_writer, PrepareEncodedArray(6, 0, /*dictionary_encoded=*/true)); + AddStructArrayOnce(format_writer, PrepareEncodedArray(4, 6, /*dictionary_encoded=*/false)); + AddStructArrayOnce(format_writer, PrepareEncodedArray(5, 10, /*dictionary_encoded=*/true)); + + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + CheckEncodedResult(file_path, /*row_count=*/15, /*null_in_dictionary=*/false); +} + +TEST_F(ParquetFormatWriterTest, TestWriteDictionaryChangingAcrossBatches) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_changing"); + std::shared_ptr out; + std::shared_ptr format_writer = CreateEncodedWriter(file_path, &out); + + // Compacting several input files puts their different dictionaries in one output row group. + // Arrow keeps only the first one and falls back to plain encoding for the rest, so the values + // have to survive that transition - and the output stops being dictionary-encoded there, which + // is the cost of forwarding an encoding rather than rebuilding one. Pinned here because it is + // what makes a compacted file bigger than one written from materialized values. + constexpr int32_t kBatchRows = 4; + for (int32_t batch = 0; batch < 3; ++batch) { + AddStructArrayOnce(format_writer, PrepareEncodedArray(kBatchRows, batch * kBatchRows, + /*dictionary_encoded=*/true, + /*null_in_dictionary=*/false, + fmt::format("batch{}_", batch))); + } + + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + + auto file = arrow::io::ReadableFile::Open(file_path, arrow_pool_.get()); + ASSERT_TRUE(file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> reader; + auto status = ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &reader); + ASSERT_TRUE(status.ok()) << status.ToString(); + const ::parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get(); + ASSERT_EQ(3 * kBatchRows, metadata->num_rows()); + auto [dictionary_pages, data_pages] = + CountDictionaryDataPages(*metadata->RowGroup(0)->ColumnChunk(0)); + ASSERT_GT(data_pages, 0); + ASSERT_LT(dictionary_pages, data_pages) << "expected a plain fallback after the second batch"; + + std::shared_ptr<::arrow::ChunkedArray> col0_array; + ASSERT_TRUE(reader->ReadColumn(0, &col0_array).ok()); + int32_t row = 0; + for (const auto& chunk : col0_array->chunks()) { + const auto& string_array = checked_pointer_cast(chunk); + ASSERT_TRUE(string_array); + for (int64_t i = 0; i < string_array->length(); ++i, ++row) { + ASSERT_EQ(fmt::format("batch{}_{}", row / kBatchRows, row % 3), + string_array->GetString(i)); + } + } + ASSERT_EQ(3 * kBatchRows, row); +} + +TEST_F(ParquetFormatWriterTest, TestWriteDictionaryWithNullsInDictionary) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_with_nulls"); + std::shared_ptr out; + std::shared_ptr format_writer = CreateEncodedWriter(file_path, &out); + + // parquet::arrow refuses a DictionaryArray whose dictionary holds nulls, so the column is + // densified rather than failing the batch. + AddStructArrayOnce(format_writer, PrepareEncodedArray(9, 0, /*dictionary_encoded=*/true, + /*null_in_dictionary=*/true)); + + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + CheckEncodedResult(file_path, /*row_count=*/9, /*null_in_dictionary=*/true); +} + +TEST_F(ParquetFormatWriterTest, TestWriteDictionaryWithNullRows) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_null_rows"); + std::shared_ptr out; + std::shared_ptr format_writer = CreateEncodedWriter(file_path, &out); + + // Nulls in the indices, not in the dictionary values: the common shape for a nullable column, + // and the one that makes the writer derive definition levels from the indices' validity + // bitmap rather than from the values it would otherwise have materialized. + std::shared_ptr indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, null, 1, 2, null, 0]") + .ValueOrDie(); + std::shared_ptr dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "c"])").ValueOrDie(); + auto dictionary_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + std::shared_ptr encoded = + arrow::DictionaryArray::FromArrays(dictionary_type, indices, dictionary).ValueOrDie(); + std::shared_ptr flat = PrepareEncodedArray(6, 0, /*dictionary_encoded=*/false); + auto struct_array = checked_pointer_cast(flat); + arrow::ArrayVector columns = {encoded, struct_array->field(1), struct_array->field(2)}; + auto batch_array = + arrow::StructArray::Make(columns, std::vector{"col1", "col2", "col3"}) + .ValueOrDie(); + AddStructArrayOnce(format_writer, batch_array); + + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + + auto file = arrow::io::ReadableFile::Open(file_path, arrow_pool_.get()); + ASSERT_TRUE(file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> reader; + auto status = ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &reader); + ASSERT_TRUE(status.ok()) << status.ToString(); + const ::parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get(); + ASSERT_EQ(6, metadata->num_rows()); + std::unique_ptr<::parquet::ColumnChunkMetaData> column_chunk = + metadata->RowGroup(0)->ColumnChunk(0); + ASSERT_TRUE(column_chunk->is_stats_set()); + ASSERT_EQ(2, column_chunk->statistics()->null_count()); + auto [dictionary_pages, data_pages] = CountDictionaryDataPages(*column_chunk); + ASSERT_GT(data_pages, 0); + ASSERT_EQ(data_pages, dictionary_pages); + + std::shared_ptr<::arrow::ChunkedArray> col0_array; + ASSERT_TRUE(reader->ReadColumn(0, &col0_array).ok()); + ASSERT_EQ(6, col0_array->length()); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), + R"(["a", null, "b", "c", null, "a"])") + .ValueOrDie(); + ASSERT_TRUE(col0_array->Equals(arrow::ChunkedArray(expected))) + << "actual=" << col0_array->ToString(); +} + +TEST_F(ParquetFormatWriterTest, TestWriteDictionaryWithDuplicateValues) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_duplicates"); + std::shared_ptr out; + std::shared_ptr format_writer = CreateEncodedWriter(file_path, &out); + + // A dictionary whose values repeat. Nothing forbids one - a DictionaryArray only bounds-checks + // its indices - but the Parquet dict encoder de-duplicates as it inserts, so its memo table + // ends up shorter than the alphabet the indices were built against. Arrow 17 notices + // (column_writer.cc, `num_entries() != dictionary->length()`) and falls back to plain rather + // than emitting a dictionary page sized from the inflated count; older forks did not, which is + // what made this a corruption rather than a size regression. Pinned because the passthrough + // relies on that fallback existing. + std::shared_ptr indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 2, 0]").ValueOrDie(); + std::shared_ptr dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "a"])").ValueOrDie(); + auto dictionary_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + std::shared_ptr encoded = + arrow::DictionaryArray::FromArrays(dictionary_type, indices, dictionary).ValueOrDie(); + std::shared_ptr flat = PrepareEncodedArray(4, 0, /*dictionary_encoded=*/false); + auto struct_array = checked_pointer_cast(flat); + arrow::ArrayVector columns = {encoded, struct_array->field(1), struct_array->field(2)}; + auto batch_array = + arrow::StructArray::Make(columns, std::vector{"col1", "col2", "col3"}) + .ValueOrDie(); + AddStructArrayOnce(format_writer, batch_array); + + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + + auto file = arrow::io::ReadableFile::Open(file_path, arrow_pool_.get()); + ASSERT_TRUE(file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> reader; + auto status = ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &reader); + ASSERT_TRUE(status.ok()) << status.ToString(); + const ::parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get(); + ASSERT_EQ(4, metadata->num_rows()); + auto [dictionary_pages, data_pages] = + CountDictionaryDataPages(*metadata->RowGroup(0)->ColumnChunk(0)); + ASSERT_GT(data_pages, 0); + ASSERT_EQ(0, dictionary_pages) << "expected the duplicate dictionary to force plain encoding"; + + // The point of the fallback: the values still come back exactly as the indices addressed them. + std::shared_ptr<::arrow::ChunkedArray> col0_array; + ASSERT_TRUE(reader->ReadColumn(0, &col0_array).ok()); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::utf8(), R"(["a", "b", "a", "a"])") + .ValueOrDie(); + ASSERT_TRUE(col0_array->Equals(arrow::ChunkedArray(expected))) + << "actual=" << col0_array->ToString(); +} + +TEST_F(ParquetFormatWriterTest, TestWriteDictionaryEmptyBatch) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_empty"); + std::shared_ptr out; + std::shared_ptr format_writer = CreateEncodedWriter(file_path, &out); + + // A rewrite can hand over an empty batch, and its layout still has to be resolved: the columns + // carry a dictionary the schema does not declare even with no rows behind it. + AddStructArrayOnce(format_writer, PrepareEncodedArray(0, 0, /*dictionary_encoded=*/true)); + AddStructArrayOnce(format_writer, PrepareEncodedArray(3, 0, /*dictionary_encoded=*/true)); + + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + CheckEncodedResult(file_path, /*row_count=*/3, /*null_in_dictionary=*/false); +} + +TEST_F(ParquetFormatWriterTest, TestGetEstimateLengthWithDictionaryBatches) { + // RollingFileWriter decides when to start a new data file from ReachTargetSize(), which is + // GetEstimateLength() against the target. A dictionary-encoded batch buffers indices rather + // than values, so the estimate is built from different bytes than it used to be; if it stopped + // tracking the file, a compaction rewrite would produce one file of unbounded size instead of + // rolling at `target-file-size`. + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_estimate_length"); + std::shared_ptr out; + // Small enough that every batch after the first opens a new buffered row group and flushes the + // previous one, so the estimate has to move for reasons the test controls. + std::shared_ptr format_writer = + CreateEncodedWriter(file_path, &out, /*max_memory_use=*/1); + + AddStructArrayOnce(format_writer, PrepareEncodedArray(64, 0, /*dictionary_encoded=*/true)); + ASSERT_OK_AND_ASSIGN(uint64_t estimate_after_first, format_writer->GetEstimateLength()); + ASSERT_GT(estimate_after_first, 0); + + ASSERT_OK_AND_ASSIGN(bool reached_tiny_target, + format_writer->ReachTargetSize(/*suggested_check=*/true, + /*target_size=*/1)); + ASSERT_TRUE(reached_tiny_target); + // Not a suggested check: the writer must not go looking at its own size at all. + ASSERT_OK_AND_ASSIGN(bool reached_unsuggested, + format_writer->ReachTargetSize(/*suggested_check=*/false, + /*target_size=*/1)); + ASSERT_FALSE(reached_unsuggested); + ASSERT_OK_AND_ASSIGN(bool reached_huge_target, + format_writer->ReachTargetSize(/*suggested_check=*/true, + /*target_size=*/1LL << 40)); + ASSERT_FALSE(reached_huge_target); + + AddStructArrayOnce(format_writer, PrepareEncodedArray(64, 64, /*dictionary_encoded=*/true)); + ASSERT_OK_AND_ASSIGN(uint64_t estimate_after_second, format_writer->GetEstimateLength()); + ASSERT_GT(estimate_after_second, estimate_after_first); + + // A flat batch after an encoded one keeps the estimate moving in the same direction, so the + // rolling decision does not depend on which encoding the rewrite happens to be forwarding. + AddStructArrayOnce(format_writer, PrepareEncodedArray(64, 128, /*dictionary_encoded=*/false)); + ASSERT_OK_AND_ASSIGN(uint64_t estimate_after_third, format_writer->GetEstimateLength()); + ASSERT_GT(estimate_after_third, estimate_after_second); + + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + ASSERT_GT(fs_->GetFileStatus(file_path).value().GetLen(), 0); +} + +TEST_F(ParquetFormatWriterTest, TestWriteDictionaryOfBinaryColumn) { + // BINARY is the other value type an `ArrowArray` layout can describe, so the writer takes a + // `dictionary(int32, binary)` batch against a plain BINARY schema exactly as it takes a STRING + // one. ParquetFileBatchReader does not currently hand one over - it forwards STRING alone, + // because the value accessors cannot read a BINARY dictionary - so this pins the writer half + // of the contract on its own, and would fail if the layout predicate were narrowed to STRING + // to enforce the reader's restriction in the wrong place. + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_binary"); + arrow::FieldVector fields = {arrow::field("b", arrow::binary())}; + ASSERT_OK_AND_ASSIGN(std::shared_ptr out, + fs_->Create(file_path, /*overwrite=*/true)); + ::parquet::WriterProperties::Builder builder; + builder.enable_dictionary(); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr format_writer, + ParquetFormatWriter::Create(out, std::make_shared(fields), builder.build(), + DEFAULT_PARQUET_WRITER_MAX_MEMORY_USE, arrow_pool_)); + + std::shared_ptr indices = + arrow::ipc::internal::json::ArrayFromJSON(arrow::int32(), "[0, 1, 0, 2]").ValueOrDie(); + std::shared_ptr dictionary = + arrow::ipc::internal::json::ArrayFromJSON(arrow::binary(), R"(["a", "bb", "ccc"])") + .ValueOrDie(); + std::shared_ptr encoded = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::binary()), + indices, dictionary) + .ValueOrDie(); + auto batch_array = + arrow::StructArray::Make({encoded}, std::vector{"b"}).ValueOrDie(); + AddStructArrayOnce(format_writer, batch_array); + + ASSERT_OK(format_writer->Flush()); + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); + + auto file = arrow::io::ReadableFile::Open(file_path, arrow_pool_.get()); + ASSERT_TRUE(file.ok()); + std::unique_ptr<::parquet::arrow::FileReader> reader; + auto status = ::parquet::arrow::OpenFile(file.ValueOrDie(), arrow_pool_.get(), &reader); + ASSERT_TRUE(status.ok()) << status.ToString(); + const ::parquet::FileMetaData* metadata = reader->parquet_reader()->metadata().get(); + ASSERT_EQ(4, metadata->num_rows()); + // The indices went to Parquet as indices: one dictionary, no plain fallback. + auto [dictionary_pages, data_pages] = + CountDictionaryDataPages(*metadata->RowGroup(0)->ColumnChunk(0)); + ASSERT_GT(data_pages, 0); + ASSERT_EQ(data_pages, dictionary_pages); + + std::shared_ptr<::arrow::ChunkedArray> column; + ASSERT_TRUE(reader->ReadColumn(0, &column).ok()); + std::shared_ptr expected = + arrow::ipc::internal::json::ArrayFromJSON(arrow::binary(), R"(["a", "bb", "a", "ccc"])") + .ValueOrDie(); + ASSERT_TRUE(column->Equals(arrow::ChunkedArray(expected))) << "actual=" << column->ToString(); +} + +TEST_F(ParquetFormatWriterTest, TestWriteDictionaryOfUnsupportedTypeIsRejected) { + std::string file_path = PathUtil::JoinPath(dir_->Str(), "dictionary_unsupported"); + std::shared_ptr out; + std::shared_ptr format_writer = CreateEncodedWriter(file_path, &out); + + // The backstop, not the behaviour a rewrite relies on: a batch layout cannot describe a + // dictionary over a non-binary-like column, and this writer only ever sees a layout, so it + // rejects rather than reinterpreting with a guessed index width. Callers that can produce such + // a column decode it while its type is still known - see + // ArrowUtils::FlattenUnresolvableDictionaries, which leaves the other columns encoded. + std::shared_ptr encoded = PrepareEncodedArray(3, 0, /*dictionary_encoded=*/false); + auto struct_array = checked_pointer_cast(encoded); + arrow::Int32Builder index_builder; + arrow::Int32Builder value_builder; + for (int32_t i = 0; i < 3; ++i) { + ASSERT_TRUE(index_builder.Append(i).ok()); + ASSERT_TRUE(value_builder.Append(7 + i).ok()); + } + std::shared_ptr indices, dictionary; + ASSERT_TRUE(index_builder.Finish(&indices).ok()); + ASSERT_TRUE(value_builder.Finish(&dictionary).ok()); + auto dictionary_int = + arrow::DictionaryArray::FromArrays(arrow::dictionary(arrow::int32(), arrow::int32()), + indices, dictionary) + .ValueOrDie(); + auto batch_array = + arrow::StructArray::Make({struct_array->field(0), dictionary_int, struct_array->field(2)}, + std::vector{"col1", "col2", "col3"}) + .ValueOrDie(); + + auto arrow_array = std::make_unique(); + ASSERT_TRUE(arrow::ExportArray(*batch_array, arrow_array.get()).ok()); + Status status = format_writer->AddBatch(arrow_array.get()); + ASSERT_TRUE(status.IsNotImplemented()) << status.ToString(); + ArrowArrayRelease(arrow_array.get()); + + ASSERT_OK(format_writer->Finish()); + ASSERT_OK(out->Flush()); + ASSERT_OK(out->Close()); +} + } // namespace paimon::parquet::test diff --git a/test/inte/append_compaction_inte_test.cpp b/test/inte/append_compaction_inte_test.cpp index ebeb23361..4031fa1b0 100644 --- a/test/inte/append_compaction_inte_test.cpp +++ b/test/inte/append_compaction_inte_test.cpp @@ -41,6 +41,8 @@ #include "paimon/file_store_commit.h" #include "paimon/file_store_write.h" #include "paimon/format/file_format_factory.h" +#include "paimon/predicate/literal.h" +#include "paimon/predicate/predicate_builder.h" #include "paimon/read_context.h" #include "paimon/result.h" #include "paimon/table/source/table_read.h" @@ -148,6 +150,34 @@ class AppendCompactionInteTest : public testing::Test, } } + // Reports the struct type of the first batch `path` yields when read through the format layer + // with `options`. The encoding a reader hands over only exists between the reader and its + // consumer - ReadResultCollector decodes it on the way out - so the raw batch is imported here + // rather than going through the table read path. + void ReadFirstBatchType(const std::shared_ptr& file_system, const std::string& path, + const std::string& file_format, + const std::map& options, + std::shared_ptr* batch_type) { + ASSERT_OK_AND_ASSIGN(auto unique_stream, file_system->Open(path)); + std::shared_ptr stream(std::move(unique_stream)); + ASSERT_OK_AND_ASSIGN(auto format, FileFormatFactory::Get(file_format, options)); + ASSERT_OK_AND_ASSIGN(auto reader_builder, format->CreateReaderBuilder(10)); + ASSERT_OK_AND_ASSIGN(auto reader, reader_builder->Build(stream)); + ASSERT_OK_AND_ASSIGN(auto c_schema, reader->GetFileSchema()); + ASSERT_OK(reader->SetReadSchema(c_schema.get(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch batch, reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(batch)); + auto& [c_array, c_array_schema] = batch; + std::shared_ptr array = + arrow::ImportArray(c_array.get(), c_array_schema.get()).ValueOrDie(); + *batch_type = array->type(); + } + + static std::string DataFilePath(const std::shared_ptr& split, size_t index = 0) { + return PathUtil::JoinPath(split->BucketPath(), split->DataFiles()[index]->file_name); + } + private: std::shared_ptr pool_; }; @@ -822,4 +852,377 @@ TEST_F(AppendCompactionInteTest, TestAppendTableCompactionWithIOException) { ASSERT_TRUE(compaction_run_complete); } +// Rewriting through the dictionary passthrough has to produce the same table as rewriting through +// materialized values, whatever encoding each input file happens to carry. The interesting part is +// the chain the unit tests cannot reach on their own: CompactRewrite hands the batch to the file +// index writer and to the format writer, and both of them recover each column's encoding from the +// batch layout after the type has been dropped by the C data interface. +// +// Both parameters set `parquet.read.enable-dictionary-passthrough` on the table, which is what the +// rewrite requires - it never enables the passthrough by itself. Parquet then forwards the +// encoding; ORC has the option vetoed because its writer cannot take a dictionary-encoded batch. +// ORC lazy decoding is on throughout, which makes the ORC reader hand over +// `dictionary(int64, large_utf8)` regardless, so the ORC parameter covers decoding a shape no +// layout can resolve. ArrowUtilsTest.TestFlattenUnresolvableDictionaries separately covers a +// vetoed writer receiving the layout-recoverable dictionary(int32, utf8) shape. +TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) { + auto file_format = GetParam(); + if (file_format != "parquet" && file_format != "orc") { + GTEST_SKIP() << file_format << " has no dictionary encoding to forward or to suppress"; + } + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + + // `s` is low-cardinality STRING and comes back encoded, so the rewrite carries an encoded and + // a materialized column at once. The other three are each excluded for their own reason: `id` + // is INT32, which the gate excludes by physical type; `b` is BINARY, which Parquet stores in + // the same BYTE_ARRAY leaf as STRING but which no value accessor can read as a dictionary; `u` + // holds a distinct value per row, the shape passthrough saves least on. The cardinality-driven + // half of the gate - a column that starts dictionary-encoded and falls back to plain partway + // through a file - needs more rows than a readable fixture holds and is covered by + // ParquetFileBatchReaderTest.TestDictionaryPassthroughSkipsFallbackToPlain instead. + arrow::FieldVector fields = { + arrow::field("id", arrow::int32()), arrow::field("s", arrow::utf8()), + arrow::field("b", arrow::binary()), arrow::field("u", arrow::utf8())}; + auto schema = arrow::schema(fields); + + std::map options = { + {Options::FILE_FORMAT, file_format}, + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, + {Options::FILE_SYSTEM, "local"}, + // Opt in explicitly: the rewrite never turns the passthrough on by itself. On the ORC + // parameter this doubles as the assertion that the option alone is not enough. + {"parquet.read.enable-dictionary-passthrough", "true"}, + {"orc.read.enable-lazy-decoding", "true"}, + // Above the distinct/total ratio of every column here, so ORC dictionary-encodes rather + // than leaving it to its own heuristic and making the assertion below data-dependent. + {"orc.dictionary-key-size-threshold", "0.9"}, + {"file-index.bitmap.columns", "s"}, + {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}, + }; + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(dir->Str(), schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/true)); + + // Each commit becomes its own input file with its own dictionary, so the rewrite has to carry + // several different dictionaries into one output column chunk. + const std::vector batches = { + R"([[1, "aa", "p", "distinct_value_1"], + [2, "bb", "q", "distinct_value_2"], + [3, null, "p", "distinct_value_3"], + [4, "aa", "q", "distinct_value_4"]])", + R"([[5, "cc", "r", "distinct_value_5"], + [6, "dd", "r", "distinct_value_6"], + [7, "cc", "s", "distinct_value_7"], + [8, "dd", "s", "distinct_value_8"]])", + R"([[9, "aa", "p", "distinct_value_9"], + [10, "ee", "t", "distinct_value_10"]])", + }; + int64_t commit_identifier = 0; + for (const std::string& data : batches) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), data, + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), commit_identifier++, + /*expected_commit_messages=*/std::nullopt)); + } + + // The rewrite must not change these, and neither must reading them back encoded. + arrow::FieldVector fields_with_row_kind = fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + auto read_type = arrow::struct_(fields_with_row_kind); + const std::string expected_rows = R"([ + [0, 1, "aa", "p", "distinct_value_1"], + [0, 2, "bb", "q", "distinct_value_2"], + [0, 3, null, "p", "distinct_value_3"], + [0, 4, "aa", "q", "distinct_value_4"], + [0, 5, "cc", "r", "distinct_value_5"], + [0, 6, "dd", "r", "distinct_value_6"], + [0, 7, "cc", "s", "distinct_value_7"], + [0, 8, "dd", "s", "distinct_value_8"], + [0, 9, "aa", "p", "distinct_value_9"], + [0, 10, "ee", "t", "distinct_value_10"] + ])"; + + { + // What the rewrite is about to read, checked on an input file with the same options the + // rewrite uses. Asserting the read type here is what makes the two directions facts of + // this test rather than assumptions: on Parquet `s` arrives as DictionaryArray while `id` + // and `b` do not; on ORC `s` arrives as `dictionary(int64, large_utf8)`, the shape no + // ArrowArray layout can resolve and the one FlattenUnresolvableDictionaries has to decode + // before the batch reaches the writer. + ASSERT_OK_AND_ASSIGN(std::vector> input_splits, + helper->NewScan(StartupMode::LatestFull(), + /*snapshot_id=*/std::nullopt)); + ASSERT_EQ(1, input_splits.size()); + auto input_split = std::dynamic_pointer_cast(input_splits[0]); + ASSERT_TRUE(input_split); + ASSERT_EQ(3, input_split->DataFiles().size()); + std::shared_ptr input_type; + ASSERT_NO_FATAL_FAILURE(ReadFirstBatchType(dir->GetFileSystem(), DataFilePath(input_split), + file_format, options, &input_type)); + ASSERT_EQ(arrow::Type::INT32, input_type->field(0)->type()->id()) << "column id"; + ASSERT_EQ(arrow::Type::DICTIONARY, input_type->field(1)->type()->id()) << "column s"; + if (file_format == "parquet") { + ASSERT_TRUE(input_type->field(1)->type()->Equals( + *arrow::dictionary(arrow::int32(), arrow::utf8()))) + << input_type->field(1)->type()->ToString(); + // `b` is dictionary-encoded in the file exactly as `s` is - same physical type, same + // pages - and is still handed over materialized, because the gate stops at STRING. + ASSERT_EQ(arrow::Type::BINARY, input_type->field(2)->type()->id()) << "column b"; + } else { + // The ORC adapter only dictionary-encodes STRING, so `b` stays materialized, and it + // widens the values: `dictionary(int64, large_utf8)` is exactly the shape whose index + // and offset widths an ArrowArray layout cannot report. + ASSERT_TRUE(input_type->field(1)->type()->Equals( + *arrow::dictionary(arrow::int64(), arrow::large_utf8()))) + << input_type->field(1)->type()->ToString(); + } + + // The same files through the whole read path rather than the raw format reader: with the + // option on, this is where an encoded batch meets the projection, the concatenation and + // the collector. It only works here - after the rewrite the output column has fallen back + // to plain, so the read at the end of this test no longer carries a dictionary at all. + ASSERT_OK_AND_ASSIGN(bool input_read_success, + helper->ReadAndCheckResult(read_type, input_splits, expected_rows)); + ASSERT_TRUE(input_read_success); + } + + ASSERT_OK(helper->write_->Compact(/*partition=*/{}, /*bucket=*/0, /*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN( + std::vector> commit_messages, + helper->write_->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); + ASSERT_FALSE(commit_messages.empty()); + ASSERT_OK(helper->commit_->Commit(commit_messages, commit_identifier)); + ASSERT_OK_AND_ASSIGN(std::optional snapshot, helper->LatestSnapshot()); + ASSERT_TRUE(snapshot); + ASSERT_EQ(Snapshot::CommitKind::Compact(), snapshot.value().GetCommitKind()); + + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + ASSERT_EQ(1, data_splits.size()); + auto data_split = std::dynamic_pointer_cast(data_splits[0]); + ASSERT_TRUE(data_split); + ASSERT_EQ(1, data_split->DataFiles().size()); + + if (file_format == "parquet") { + // What the passthrough costs, on the file it just produced. Three input files brought + // three different dictionaries into one output row group, so the writer kept the first and + // wrote the rest plain; the gate then declines the output column, which is what makes it + // ineligible for the next compaction round and bigger than a rewrite from materialized + // values. TestAppendTableCompactionDictionaryPassthroughDefaultOff is the other side. + std::shared_ptr output_type; + ASSERT_NO_FATAL_FAILURE(ReadFirstBatchType(dir->GetFileSystem(), DataFilePath(data_split), + file_format, options, &output_type)); + ASSERT_NE(arrow::Type::DICTIONARY, output_type->field(1)->type()->id()) + << "column s: " << output_type->field(1)->type()->ToString(); + } + + ASSERT_OK_AND_ASSIGN(bool success, + helper->ReadAndCheckResult(read_type, data_splits, expected_rows)); + ASSERT_TRUE(success); + + // The bitmap index on `s` is built from the rewritten batch, which reaches the index writer + // still encoded when the passthrough is on. Reading through the index is what proves it was + // decoded against the right values rather than against its indices. + std::string indexed_value = "cc"; + auto predicate = PredicateBuilder::Equal( + /*field_index=*/1, /*field_name=*/"s", FieldType::STRING, + Literal(FieldType::STRING, indexed_value.data(), indexed_value.size())); + ReadContextBuilder read_context_builder(PathUtil::JoinPath(dir->Str(), "foo.db/bar")); + read_context_builder.SetOptions(options).SetPredicate(predicate); + ASSERT_OK_AND_ASSIGN(auto read_context, read_context_builder.Finish()); + ASSERT_OK_AND_ASSIGN(auto table_read, TableRead::Create(std::move(read_context))); + ASSERT_OK_AND_ASSIGN(auto batch_reader, table_read->CreateReader(data_splits)); + ASSERT_OK_AND_ASSIGN(auto filtered, ReadResultCollector::CollectResult(batch_reader.get())); + auto expected = std::make_shared( + arrow::ipc::internal::json::ArrayFromJSON(read_type, R"([ + [0, 5, "cc", "r", "distinct_value_5"], + [0, 7, "cc", "s", "distinct_value_7"] + ])") + .ValueOrDie()); + ASSERT_TRUE(expected->Equals(filtered)) << "actual=" << filtered->ToString(); +} + +// The same rewrite on a table that says nothing about the passthrough, which is the default and +// the shape of every table that has not opted in. It has to land on the same rows as +// TestAppendTableCompactionDictionaryPassthrough, so the option is a performance knob and never a +// correctness one - and on a *better* encoded output, which is why it is not turned on by default. +TEST_F(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthroughDefaultOff) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("s", arrow::utf8())}; + auto schema = arrow::schema(fields); + // No `parquet.read.enable-dictionary-passthrough` here on purpose: absence is what is under + // test. An explicit `false` reaches the same code, since the reader resolves the option + // against the same default. + std::map options = { + {Options::FILE_FORMAT, "parquet"}, + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, + {Options::FILE_SYSTEM, "local"}, + }; + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(dir->Str(), schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/true)); + + // Three files, which is what full compaction needs before it rewrites anything, and three + // different dictionaries for the writer to reconcile. + const std::vector batches = { + R"([[1, "aa"], [2, "bb"]])", R"([[3, "cc"], [4, "aa"]])", R"([[5, "dd"], [6, "cc"]])"}; + int64_t commit_identifier = 0; + for (const std::string& data : batches) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), data, + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), commit_identifier++, + /*expected_commit_messages=*/std::nullopt)); + } + + ASSERT_OK(helper->write_->Compact(/*partition=*/{}, /*bucket=*/0, /*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN( + std::vector> commit_messages, + helper->write_->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); + ASSERT_FALSE(commit_messages.empty()); + ASSERT_OK(helper->commit_->Commit(commit_messages, commit_identifier)); + + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + ASSERT_EQ(1, data_splits.size()); + auto data_split = std::dynamic_pointer_cast(data_splits[0]); + ASSERT_TRUE(data_split); + ASSERT_EQ(1, data_split->DataFiles().size()); + + // The point of leaving the passthrough off: the writer built one dictionary covering the whole + // row group, so `s` is dictionary-encoded end to end and the gate accepts it - unlike the same + // rewrite with the passthrough on, which leaves the column partly plain. Read with the option + // on, since that is what makes the encoding visible; it does not change what is in the file. + std::map passthrough_options = options; + passthrough_options["parquet.read.enable-dictionary-passthrough"] = "true"; + std::shared_ptr output_type; + ASSERT_NO_FATAL_FAILURE(ReadFirstBatchType(dir->GetFileSystem(), DataFilePath(data_split), + "parquet", passthrough_options, &output_type)); + ASSERT_EQ(arrow::Type::DICTIONARY, output_type->field(1)->type()->id()) + << "column s: " << output_type->field(1)->type()->ToString(); + + arrow::FieldVector fields_with_row_kind = fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + ASSERT_OK_AND_ASSIGN(bool success, helper->ReadAndCheckResult( + arrow::struct_(fields_with_row_kind), data_splits, R"([ + [0, 1, "aa"], + [0, 2, "bb"], + [0, 3, "cc"], + [0, 4, "aa"], + [0, 5, "dd"], + [0, 6, "cc"] + ])")); + ASSERT_TRUE(success); +} + +// The veto in GetDictionaryPassthroughVetoReason() that no other test reaches: a shredding writer +// reshapes every batch against a fixed physical schema, so the rewrite forces the read option off +// even though the table asks for it. The format veto is covered by the ORC parameter of +// TestAppendTableCompactionDictionaryPassthrough; the `parquet.enable-dictionary` one is left +// uncovered on purpose, since a writer with dictionaries disabled densifies whatever it is handed +// and the rewritten file is identical either way. +// +// `s` is what makes this one visible: it reaches the rewrite dictionary-encoded - asserted on an +// input file first, so the test cannot pass merely because there was nothing to forward - and comes +// back dictionary-encoded end to end, which only a rewrite from materialized values produces. Had +// the option been honoured it would have gone plain after the first input file's dictionary, the +// way it does in TestAppendTableCompactionDictionaryPassthrough. +TEST_F(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthroughVetoedByShredding) { + auto dir = UniqueTestDirectory::Create(); + ASSERT_TRUE(dir); + auto map_type = arrow::map(arrow::utf8(), arrow::int64()); + arrow::FieldVector fields = {arrow::field("id", arrow::int32()), + arrow::field("s", arrow::utf8()), arrow::field("tags", map_type)}; + auto schema = arrow::schema(fields); + std::map options = { + {Options::FILE_FORMAT, "parquet"}, + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, + {Options::FILE_SYSTEM, "local"}, + {"parquet.read.enable-dictionary-passthrough", "true"}, + // The shredding plan, which is what makes the rewrite veto the option above. + {"fields.tags.map.storage-layout", "shared-shredding"}, + {"fields.tags.map.shared-shredding.max-columns", "64"}, + }; + ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(dir->Str(), schema, /*partition_keys=*/{}, + /*primary_keys=*/{}, options, + /*is_streaming_mode=*/true)); + + // Three files, each with its own `s` dictionary, which is what full compaction needs before it + // rewrites anything. + const std::vector batches = { + R"([[1, "aa", [["a", 10]]], [2, "bb", [["b", 20]]]])", + R"([[3, "cc", [["a", 30]]], [4, "aa", null]])", + R"([[5, "dd", [["c", 40]]], [6, "cc", [["a", 50], ["d", 60]]]])", + }; + int64_t commit_identifier = 0; + for (const std::string& data : batches) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr batch, + TestHelper::MakeRecordBatch(arrow::struct_(fields), data, + /*partition_map=*/{}, /*bucket=*/0, {})); + ASSERT_OK(helper->WriteAndCommit(std::move(batch), commit_identifier++, + /*expected_commit_messages=*/std::nullopt)); + } + + { + // Without this the veto could not be told apart from an input the gate declines anyway. + ASSERT_OK_AND_ASSIGN(std::vector> input_splits, + helper->NewScan(StartupMode::LatestFull(), + /*snapshot_id=*/std::nullopt)); + ASSERT_EQ(1, input_splits.size()); + auto input_split = std::dynamic_pointer_cast(input_splits[0]); + ASSERT_TRUE(input_split); + ASSERT_EQ(3, input_split->DataFiles().size()); + std::shared_ptr input_type; + ASSERT_NO_FATAL_FAILURE(ReadFirstBatchType(dir->GetFileSystem(), DataFilePath(input_split), + "parquet", options, &input_type)); + ASSERT_EQ("s", input_type->field(1)->name()); + ASSERT_EQ(arrow::Type::DICTIONARY, input_type->field(1)->type()->id()) + << "column s: " << input_type->field(1)->type()->ToString(); + } + + ASSERT_OK(helper->write_->Compact(/*partition=*/{}, /*bucket=*/0, /*full_compaction=*/true)); + ASSERT_OK_AND_ASSIGN( + std::vector> commit_messages, + helper->write_->PrepareCommit(/*wait_compaction=*/true, commit_identifier)); + ASSERT_FALSE(commit_messages.empty()); + ASSERT_OK(helper->commit_->Commit(commit_messages, commit_identifier)); + + ASSERT_OK_AND_ASSIGN(std::vector> data_splits, + helper->NewScan(StartupMode::LatestFull(), /*snapshot_id=*/std::nullopt)); + ASSERT_EQ(1, data_splits.size()); + auto data_split = std::dynamic_pointer_cast(data_splits[0]); + ASSERT_TRUE(data_split); + ASSERT_EQ(1, data_split->DataFiles().size()); + + std::shared_ptr output_type; + ASSERT_NO_FATAL_FAILURE(ReadFirstBatchType(dir->GetFileSystem(), DataFilePath(data_split), + "parquet", options, &output_type)); + ASSERT_EQ("s", output_type->field(1)->name()); + ASSERT_EQ(arrow::Type::DICTIONARY, output_type->field(1)->type()->id()) + << "column s: " << output_type->field(1)->type()->ToString(); + + arrow::FieldVector fields_with_row_kind = fields; + fields_with_row_kind.insert(fields_with_row_kind.begin(), + arrow::field("_VALUE_KIND", arrow::int8())); + ASSERT_OK_AND_ASSIGN(bool success, helper->ReadAndCheckResult( + arrow::struct_(fields_with_row_kind), data_splits, R"([ + [0, 1, "aa", [["a", 10]]], + [0, 2, "bb", [["b", 20]]], + [0, 3, "cc", [["a", 30]]], + [0, 4, "aa", null], + [0, 5, "dd", [["c", 40]]], + [0, 6, "cc", [["a", 50], ["d", 60]]] + ])")); + ASSERT_TRUE(success); +} + } // namespace paimon::test