From 450eed27ff02b6a485925ced7b4cc6fb6964d3bd Mon Sep 17 00:00:00 2001 From: Nicholas Jiang Date: Fri, 28 Aug 2026 14:24:47 +0800 Subject: [PATCH 1/3] perf(parquet): forward dictionary encoding through append compaction rewrite An append-only compaction rewrite copies rows into the new file without inspecting any value, so expanding a dictionary-encoded Parquet column on read and hashing it again on write is work neither side needs. Forward the encoding instead. Reader: `parquet.read.enable-dictionary-passthrough`, off by default, makes ParquetFileBatchReader request `set_read_dictionary` for non-nested STRING/BINARY columns whose every data page, in every row group of the file, is dictionary-encoded. A dictionary page alone cannot be the signal - a column that outgrew its page limit still carries the page it had already emitted - so the gate reads the encoding statistics. Writer: ParquetFormatWriter recovers each batch's encoding from its layout, because exporting through the Arrow C data interface drops the type. A layout pins down neither the index nor the offset width, so only `dictionary(int32, utf8|binary)` is recoverable; CompactRewrite decodes anything else per column while the type is still known - the ORC reader's `dictionary(int64, large_utf8)` under lazy decoding, dictionaries below the top level - and leaves the rest encoded. A dictionary holding nulls in its values is flattened at the writer, the one shape parquet::arrow rejects outright. Compaction opts in only when the output is Parquet, `parquet.enable-dictionary` is on and no shredding plan is active; anything else forces the read option off. A file index on a forwarded column materializes that column alone. Note that a Parquet column chunk carries one dictionary, so when the input files supply different ones the output keeps the first and falls back to plain for the rest of the row group. The rewritten data is unchanged, but the output file may be larger than one written from materialized values. --- benchmark/parquet_format_benchmark.cpp | 97 ++++ docs/source/user_guide/compaction.rst | 34 ++ .../common/reader/reader_utils_test.cpp | 47 ++ src/paimon/common/utils/arrow/arrow_utils.cpp | 143 ++++++ src/paimon/common/utils/arrow/arrow_utils.h | 80 ++++ .../common/utils/arrow/arrow_utils_test.cpp | 247 ++++++++++ src/paimon/core/io/data_file_index_writer.cpp | 40 +- .../core/io/data_file_index_writer_test.cpp | 41 ++ src/paimon/core/io/data_file_writer_base.h | 16 +- .../append_only_file_store_write.cpp | 66 ++- .../operation/append_only_file_store_write.h | 20 +- .../parquet/parquet_file_batch_reader.cpp | 156 ++++++- .../parquet/parquet_file_batch_reader.h | 23 +- .../parquet_file_batch_reader_test.cpp | 266 +++++++++++ .../format/parquet/parquet_format_defs.h | 12 + .../format/parquet/parquet_format_writer.cpp | 55 ++- .../format/parquet/parquet_format_writer.h | 26 ++ .../parquet/parquet_format_writer_test.cpp | 430 ++++++++++++++++++ test/inte/append_compaction_inte_test.cpp | 237 ++++++++++ 19 files changed, 2012 insertions(+), 24 deletions(-) diff --git a/benchmark/parquet_format_benchmark.cpp b/benchmark/parquet_format_benchmark.cpp index e8d3e0026..b4e93444a 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)}); } @@ -422,6 +432,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 +788,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 +973,21 @@ 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. +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); +} + // 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 +1255,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 +1346,16 @@ 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(); BENCHMARK(BM_ParquetWrite_DictionaryInt32) ->ArgName("cardinality") ->Arg(10) @@ -1402,6 +1486,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/docs/source/user_guide/compaction.rst b/docs/source/user_guide/compaction.rst index 093313555..0a0b561ca 100644 --- a/docs/source/user_guide/compaction.rst +++ b/docs/source/user_guide/compaction.rst @@ -88,6 +88,40 @@ 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``/``BINARY`` columns +benefit most. Primary-key compaction merges rows and is not covered. + +This applies automatically. Eligibility is decided per input file: a non-nested +``STRING``/``BINARY`` 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. Passthrough is also skipped 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. + +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. + +Passthrough changes what the rewrite costs, not what it produces, with one +exception worth knowing: a Parquet column chunk can only carry one dictionary, +so when the input files supply different dictionaries the output column keeps +the first and falls back to plain encoding for the rest of the row group. The +rewritten data is unchanged either way, but the output file may be larger than a +rewrite that rebuilt a single dictionary from materialized values. Set +``parquet.read.enable-dictionary-passthrough`` to ``false`` on the table to turn +the optimization off and always rebuild. + 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..a69b677b9 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -24,6 +24,9 @@ #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/compute/exec.h" #include "arrow/type_traits.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" @@ -38,6 +41,55 @@ 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::IsParquetDictionaryValueType(). +bool IsResolvableDictionary(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::IsParquetDictionaryValueType(*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,95 @@ Result ArrowUtils::GetCompressionType(const std::strin return compression_type; } +bool ArrowUtils::IsParquetDictionaryValueType(const arrow::DataType& type) { + return type.id() == arrow::Type::STRING || type.id() == arrow::Type::BINARY; +} + +Result> ArrowUtils::ResolveParquetDictionaryStructType( + 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 (!IsParquetDictionaryValueType(*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) { + 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); + arrow::compute::ExecContext exec_context(pool); + 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 (IsResolvableDictionary(*field->type()) || !HasDictionary(*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_FROM_ARROW( + arrow::Datum decoded, + arrow::compute::Cast(arrow::MakeArray(data->child_data[i]), logical_field->type(), + arrow::compute::CastOptions::Safe(), &exec_context)); + data->child_data[i] = decoded.array(); + 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..e2e6c3ed6 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,84 @@ class PAIMON_EXPORT ArrowUtils { /// Handles "none" and empty string by mapping them to "uncompressed". static Result GetCompressionType(const std::string& compression); + /// Whether a column of `type` may be carried dictionary-encoded across the Arrow C data + /// interface, which drops the type and leaves only the layout behind. + /// + /// A layout pins down neither the index width nor the offset width, so the only encoding worth + /// carrying is the one a single known producer emits: `dictionary(int32(), utf8()|binary())`, + /// which is what Arrow's Parquet reader produces for + /// `ArrowReaderProperties::set_read_dictionary`. `LARGE_STRING` is deliberately excluded even + /// though it is binary-like: the ORC reader widens strings to + /// `dictionary(int64(), large_utf8())` under lazy decoding, and 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 + /// ResolveParquetDictionaryStructType() for where the index width becomes a caller contract. + /// + /// This is the single definition shared by the reader that decides which columns to request + /// encoded and by the writer that has to recognise them again on the other side. + /// + /// @param type The column's value type, not its dictionary type. + /// @return True when `dictionary(int32(), type)` round-trips through an `ArrowArray`. + static bool IsParquetDictionaryValueType(const arrow::DataType& type); + + /// Recovers the struct type of a batch that Arrow's Parquet reader produced with + /// `set_read_dictionary` enabled: `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 not inferred, it is assumed, and that assumption is only valid + /// for Arrow's Parquet reader. **The value type check does not make it safe for anything + /// else**: it rejects `dictionary(int64(), large_utf8())`, which is 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 this is a contract, not a check, and it binds the code that *produces* the batch rather + /// than the two places that call this. A producer must either be handing on a batch that came + /// straight from Arrow's Parquet reader, or must run FlattenUnresolvableDictionaries() while + /// the type is still known - that one does test the index width, and decodes every column this + /// cannot resolve while leaving the rest encoded. + /// `AppendOnlyFileStoreWrite::CompactRewrite` is today's only production path that can hand + /// over a batch whose dictionaries the schema does not declare, and it takes the second route. + /// The callers themselves - `ParquetFormatWriter::ResolveBatchSchema` and + /// `DataFileWriterBase::AddFileIndexBatch` - are downstream of it and see only the layout. + /// + /// The value-type rejection and the rejection of a dictionary below the top level narrow the + /// blast radius; they do not close it. Closing it needs the real `ArrowSchema` to reach the + /// writer, which the `FormatWriter::AddBatch(ArrowArray*)` signature currently drops. + /// + /// 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> ResolveParquetDictionaryStructType( + const std::shared_ptr& logical_type, const ::ArrowArray* batch); + + /// Returns `batch` with every top-level column that ResolveParquetDictionaryStructType() could + /// not resolve decoded to the type its field carries in `logical_type`. A column it can + /// resolve stays 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. + /// + /// This is the counterpart of the restriction above: exporting an array through the C data + /// interface drops its type, so a column whose encoding does not survive that round trip 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. + /// @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); + 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..bc2266115 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,245 @@ TEST(ArrowUtilsTest, TestGetCompressionType) { } } +TEST(ArrowUtilsTest, TestResolveParquetDictionaryStructType) { + // 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::ResolveParquetDictionaryStructType(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 describable_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + // What the ORC reader hands over for a dictionary-encoded string column under lazy decoding. + auto undescribable_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 describable = + arrow::DictionaryArray::FromArrays( + describable_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 undescribable = + arrow::DictionaryArray::FromArrays( + undescribable_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({describable, undescribable, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries(batch, logical_type, pool)); + 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(*describable)); + } + { + // A dictionary below the top level is undescribable too, so the whole column is decoded. + auto nested = + arrow::StructArray::Make({undescribable}, 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)); + 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. + // ResolveParquetDictionaryStructType() 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. This is the + // single reason CompactRewrite has to run this before exporting, rather than relying on + // the writer's own rejection. + 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({describable, wide_indices, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries(batch, logical_type, pool)); + 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({describable, describable, ints}, + std::vector{"s", "o", "i"}) + .ValueOrDie()); + ASSERT_OK_AND_ASSIGN( + std::shared_ptr flattened, + ArrowUtils::FlattenUnresolvableDictionaries(batch, logical_type, pool)); + 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({describable, undescribable, 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)); + 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(); + } +} + } // namespace paimon::test diff --git a/src/paimon/core/io/data_file_index_writer.cpp b/src/paimon/core/io/data_file_index_writer.cpp index 5c97bb3da..74c964909 100644 --- a/src/paimon/core/io/data_file_index_writer.cpp +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -20,15 +20,19 @@ #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 "arrow/compute/exec.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" @@ -100,10 +104,40 @@ Status DataFileIndexWriter::AddBatch(const std::shared_ptr& if (finished_) { return Status::Invalid("Data file index writer has already finished"); } + // Buffers allocated through the adaptor keep a raw pointer to it, so it has to outlive every + // array decoded below. Built on first use, since most batches decode nothing. + std::unique_ptr arrow_pool; + // 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_); + } + arrow::compute::ExecContext exec_context(arrow_pool.get()); + PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( + arrow::Datum decoded, + arrow::compute::Cast(column, entry.field->type(), + arrow::compute::CastOptions::Safe(), &exec_context)); + column = decoded.make_array(); + 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_test.cpp b/src/paimon/core/io/data_file_index_writer_test.cpp index 1e2c9593f..db2306167 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,44 @@ 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"}})); + + 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::utf8(), R"(["a", "b", "c"])").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(), "[10, 20, 30, 40]").ValueOrDie(); + auto batch = checked_pointer_cast( + arrow::StructArray::Make({encoded, values}, std::vector{"f0", "f1"}) + .ValueOrDie()); + + ASSERT_OK(writer->AddBatch(batch)); + 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()); + ASSERT_OK_AND_ASSIGN(auto single_row_result, + bitmap_readers[0]->VisitEqual(Literal(FieldType::STRING, "b", 1))); + ASSERT_EQ("{1}", single_row_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..bb6b01fd8 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::ResolveParquetDictionaryStructType(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..7ce2f0a0b 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" @@ -55,10 +57,12 @@ #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/executor.h" +#include "paimon/format/parquet/parquet_format_defs.h" #include "paimon/logging.h" #include "paimon/read_context.h" #include "paimon/realtime/realtime_context.h" #include "paimon/result.h" +#include "parquet/properties.h" namespace arrow { class Schema; } // namespace arrow @@ -146,13 +150,24 @@ 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(bool dictionary_passthrough, CanUseDictionaryPassthrough(plan_factory)); + PAIMON_ASSIGN_OR_RAISE( + std::unique_ptr reader, + CreateFilesReader(partition, bucket, dv_factory, to_compact, dictionary_passthrough)); 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 +203,13 @@ 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 batch layout alone. Decode here, while the type is still known, whatever that + // recovery cannot describe - an ORC reader under lazy decoding hands over + // `dictionary(int64, large_utf8)`, which a layout says nothing about. Only those columns + // pay for it; a Parquet passthrough column stays encoded. + PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::FlattenUnresolvableDictionaries( + struct_array, logical_type, arrow_pool.get())); PAIMON_RETURN_NOT_OK_FROM_ARROW( arrow::ExportArray(*struct_array, c_array.get(), c_schema.get())); ArrowSchemaRelease(c_schema.get()); @@ -264,10 +286,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 +299,25 @@ Result AppendOnlyFileStoreWrite::GetDat data_file_path_factory, pool_); } +Result AppendOnlyFileStoreWrite::CanUseDictionaryPassthrough( + const std::shared_ptr& plan_factory) const { + std::shared_ptr file_format = options_.GetFileFormat(); + if (!file_format || file_format->Identifier() != "parquet") { + return false; + } + PAIMON_ASSIGN_OR_RAISE( + bool enable_dictionary, + OptionsUtils::GetValueFromMap(options_.ToMap(), parquet::PARQUET_ENABLE_DICTIONARY, + ::parquet::DEFAULT_IS_DICTIONARY_ENABLED)); + if (!enable_dictionary) { + return false; + } + return plan_factory == nullptr; +} + Result> AppendOnlyFileStoreWrite::CreateFilesReader( const BinaryRow& partition, int32_t bucket, DeletionVector::Factory dv_factory, - const std::vector>& files) const { + const std::vector>& files, bool dictionary_passthrough) const { ReadContextBuilder context_builder(root_path_); context_builder.SetOptions(options_.ToMap()) .WithFileSystem(options_.GetFileSystem()) @@ -290,6 +327,17 @@ Result> AppendOnlyFileStoreWrite::CreateFilesReader .WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, context_builder.Finish()); std::map options = options_.ToMap(); + // CompactRewrite copies batches into the rewritten file without looking at any value, so a + // column the input files already store dictionary-encoded can keep that encoding instead of + // being expanded here and hashed again by the writer. + if (dictionary_passthrough) { + // `emplace` so an explicit table option can still turn it off. + options.emplace(parquet::PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH, "true"); + } else { + // Not negotiable the other way: a writer that cannot take a dictionary-encoded batch must + // not receive one because the table happens to set the read option. + options[parquet::PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH] = "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..48dae542b 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,30 @@ 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 + /// `CanUseDictionaryPassthrough` 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; Result> CreateFilesReader( const BinaryRow& partition, int32_t bucket, DeletionVector::Factory dv_factory, - const std::vector>& files) const; + const std::vector>& files, bool dictionary_passthrough) const; + + /// Whether `CompactRewrite` may forward the dictionary encoding of its input files instead of + /// expanding every value. Requires 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. + Result CanUseDictionaryPassthrough( + 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..aacebadbe 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,96 @@ 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. + if (schema->Column(i)->physical_type() == ::parquet::Type::BYTE_ARRAY && + 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::IsParquetDictionaryValueType(*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 +295,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 +313,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 +340,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 +364,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 +443,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..95015df87 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.h +++ b/src/paimon/format/parquet/parquet_file_batch_reader.h @@ -162,13 +162,30 @@ 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 BYTE_ARRAY + /// columns whose every data page, in every row group, is dictionary-encoded. + 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 +279,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..41093eab8 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,265 @@ 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 every read outside the compaction rewrite reaches this reader: the + // default has to be off, or an ordinary scan would start emitting dictionary batches at + // consumers that do not unwrap them. + 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, 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..dc721d7a2 100644 --- a/src/paimon/format/parquet/parquet_format_defs.h +++ b/src/paimon/format/parquet/parquet_format_defs.h @@ -102,12 +102,24 @@ 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/BINARY columns as Arrow DictionaryArray instead of one copy of +// the value per row. Restricted to non-nested leaf columns whose every data page is already +// dictionary-encoded, so the reader only ever hands on a dictionary the file itself has. +// +// Off by default because it only pays off when the consumer forwards the batch without inspecting +// values, which is why the append compaction rewrite is the one caller that opts in. Value +// accessors have to unwrap DictionaryArray to read such a column; `ColumnarUtils::GetView` does, +// but that is not true of every accessor, so a new consumer has to be checked before enabling it. +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..c2d35d714 100644 --- a/src/paimon/format/parquet/parquet_format_writer.cpp +++ b/src/paimon/format/parquet/parquet_format_writer.cpp @@ -23,14 +23,19 @@ #include #include +#include "arrow/array/array_dict.h" #include "arrow/c/bridge.h" +#include "arrow/compute/api.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/format/parquet/parquet_format_defs.h" #include "parquet/arrow/writer.h" #include "parquet/properties.h" @@ -64,8 +69,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 +81,51 @@ Status ParquetFormatWriter::AddBatch(ArrowArray* batch) { return Status::OK(); } +Result> ParquetFormatWriter::ResolveBatchSchema( + const ::ArrowArray* batch) { + PAIMON_ASSIGN_OR_RAISE( + std::shared_ptr batch_type, + ArrowUtils::ResolveParquetDictionaryStructType(logical_struct_type_, batch)); + if (batch_type == logical_struct_type_) { + return schema_; + } + if (dictionary_batch_type_ == nullptr || !dictionary_batch_type_->Equals(*batch_type)) { + dictionary_batch_type_ = batch_type; + dictionary_batch_schema_ = arrow::schema(batch_type->fields(), schema_->metadata()); + } + return dictionary_batch_schema_; +} + +Result> ParquetFormatWriter::FlattenUnwritableDictionaries( + const std::shared_ptr& record_batch) const { + arrow::ArrayVector columns; + arrow::FieldVector fields; + arrow::compute::ExecContext exec_context(pool_.get()); + 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_FROM_ARROW( + arrow::Datum flattened, + arrow::compute::Cast(column, dictionary_type.value_type(), + arrow::compute::CastOptions::Safe(), &exec_context)); + columns[i] = flattened.make_array(); + 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 +171,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..a0f634397 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,34 @@ 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); + + /// 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_; + // Last dictionary-encoded batch type and its schema, so a run of identically encoded batches + // builds the import schema only once. + std::shared_ptr dictionary_batch_type_; + std::shared_ptr dictionary_batch_schema_; 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..f18c94d34 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,303 @@ 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, 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..0e370e7da 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" @@ -822,4 +824,239 @@ 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. +// +// Parameterised over the two formats that have an encoding to forward or to suppress: Parquet +// turns the passthrough on, ORC forces it off 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)` - a shape no layout can resolve, so it exercises the +// decode-at-the-source path rather than the passthrough. +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` and `b` are low-cardinality and come back encoded; `id` is INT32, which the gate excludes + // by physical type, so the rewrite carries both kinds of column at once. `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"}, + {"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)); + } + + { + // 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` and `b` arrive as DictionaryArray and + // `id` does not, because the gate only considers BYTE_ARRAY leaves; 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::string input_path = + PathUtil::JoinPath(input_split->BucketPath(), input_split->DataFiles()[0]->file_name); + ASSERT_OK_AND_ASSIGN(auto unique_input_stream, dir->GetFileSystem()->Open(input_path)); + std::shared_ptr input_stream(std::move(unique_input_stream)); + + std::map passthrough_options = options; + passthrough_options["parquet.read.enable-dictionary-passthrough"] = "true"; + ASSERT_OK_AND_ASSIGN(auto input_file_format, + FileFormatFactory::Get(file_format, passthrough_options)); + ASSERT_OK_AND_ASSIGN(auto input_reader_builder, input_file_format->CreateReaderBuilder(10)); + ASSERT_OK_AND_ASSIGN(auto input_reader, input_reader_builder->Build(input_stream)); + ASSERT_OK_AND_ASSIGN(auto c_input_schema, input_reader->GetFileSchema()); + ASSERT_OK(input_reader->SetReadSchema(c_input_schema.get(), /*predicate=*/nullptr, + /*selection_bitmap=*/std::nullopt)); + // Read one batch directly rather than through ReadResultCollector, which decodes + // dictionaries on the way out and would hide the very thing being asserted. + ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch input_batch, input_reader->NextBatch()); + ASSERT_FALSE(BatchReader::IsEofBatch(input_batch)); + auto& [input_c_array, input_c_schema] = input_batch; + std::shared_ptr input_array = + arrow::ImportArray(input_c_array.get(), input_c_schema.get()).ValueOrDie(); + std::shared_ptr input_type = input_array->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_EQ(arrow::Type::DICTIONARY, input_type->field(2)->type()->id()) << "column b"; + ASSERT_TRUE(input_type->field(1)->type()->Equals( + *arrow::dictionary(arrow::int32(), arrow::utf8()))) + << input_type->field(1)->type()->ToString(); + } 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(); + } + } + + 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()); + + 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); + ASSERT_OK_AND_ASSIGN(bool success, helper->ReadAndCheckResult(read_type, data_splits, 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"] + ])")); + 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 with the passthrough explicitly disabled has to land on the same table, so the +// kill switch is a performance knob and never a correctness one. +TEST_F(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthroughDisabled) { + 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); + std::map options = { + {Options::FILE_FORMAT, "parquet"}, + {Options::BUCKET, "1"}, + {Options::BUCKET_KEY, "id"}, + {Options::FILE_SYSTEM, "local"}, + {"parquet.read.enable-dictionary-passthrough", "false"}, + }; + 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)); + 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); +} + } // namespace paimon::test From dc2c7598503b9df92a6d565fb28814f482fd6a97 Mon Sep 17 00:00:00 2001 From: Nicholas Jiang Date: Mon, 31 Aug 2026 15:27:19 +0800 Subject: [PATCH 2/3] perf(parquet): make dictionary passthrough opt-in and restrict it to STRING Addresses the review on #257. The compaction rewrite no longer turns `parquet.read.enable-dictionary-passthrough` on. It is a table option, off by default, and the rewrite only vetoes it - when the output is not Parquet, when `parquet.enable-dictionary` is false, or when a shredding writer is active - logging the reason at DEBUG for a table that did ask for it. Forwarding trades compaction CPU for output size, and a Parquet column chunk carries one dictionary, so a multi-file rewrite can keep the first input file's and write the rest of the row group plain; whether that trade is worth taking depends on the data, so it is left to the table. Restrict the reader to STRING. Parquet stores BINARY in the same BYTE_ARRAY leaf and dictionary-encodes it the same way, but no value accessor here can read a `dictionary(int32, binary)`: ColumnarUtils::GetView() asserts on it and returns an empty view in a release build, and LiteralConverter rejects it. Since the option applies to every read of the table, the gate gets that restriction, not the layout predicate - ArrowUtils::IsDictionaryLayoutRecoverableValueType() still accepts `utf8|binary`, which is what the writer can recover, so the writer keeps the capability for a future producer. Also from the review: use CastingUtils::Cast() instead of arrow::compute::Cast(); hold the Arrow pool adaptor in DataFileIndexWriter as a member declared before the index writers, so it outlives every buffer allocated through it; drop the batch schema cache in ParquetFormatWriter::ResolveBatchSchema(); read the two Parquet option names in core as local constants rather than through the format layer's headers. The changing-dictionary benchmark now rotates one alphabet and shifts the indices back by the same amount, so it writes the same logical column as its pair and the delta is the fallback rather than a difference in the data. --- benchmark/parquet_format_benchmark.cpp | 56 +++- benchmark/parquet_format_benchmark_test.cpp | 202 +++++++++++-- docs/source/examples/benchmark.rst | 19 ++ docs/source/user_guide/compaction.rst | 78 +++-- src/paimon/common/utils/arrow/arrow_utils.cpp | 41 ++- src/paimon/common/utils/arrow/arrow_utils.h | 50 +--- src/paimon/core/casting/casting_utils.h | 3 +- src/paimon/core/io/data_file_index_writer.cpp | 21 +- src/paimon/core/io/data_file_index_writer.h | 11 + .../core/io/data_file_index_writer_test.cpp | 48 +-- .../append_only_file_store_write.cpp | 72 +++-- .../operation/append_only_file_store_write.h | 21 +- .../parquet/parquet_file_batch_reader.cpp | 13 +- .../parquet/parquet_file_batch_reader.h | 5 +- .../parquet_file_batch_reader_test.cpp | 43 ++- .../format/parquet/parquet_format_defs.h | 15 +- .../format/parquet/parquet_format_writer.cpp | 22 +- .../format/parquet/parquet_format_writer.h | 6 +- .../parquet/parquet_format_writer_test.cpp | 57 ++++ test/inte/append_compaction_inte_test.cpp | 275 ++++++++++++++---- 20 files changed, 811 insertions(+), 247 deletions(-) diff --git a/benchmark/parquet_format_benchmark.cpp b/benchmark/parquet_format_benchmark.cpp index b4e93444a..3823ac3bd 100644 --- a/benchmark/parquet_format_benchmark.cpp +++ b/benchmark/parquet_format_benchmark.cpp @@ -275,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, @@ -283,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. @@ -979,6 +1003,10 @@ void BM_ParquetWrite_DictionaryString(::benchmark::State& state) { // 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(), @@ -988,6 +1016,24 @@ void BM_ParquetWrite_DictionaryStringIntoStringSchema(::benchmark::State& state) 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 @@ -1356,6 +1402,14 @@ BENCHMARK(BM_ParquetWrite_DictionaryStringIntoStringSchema) ->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) 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 0a0b561ca..44419a245 100644 --- a/docs/source/user_guide/compaction.rst +++ b/docs/source/user_guide/compaction.rst @@ -95,32 +95,66 @@ 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``/``BINARY`` columns -benefit most. Primary-key compaction merges rows and is not covered. - -This applies automatically. Eligibility is decided per input file: a non-nested -``STRING``/``BINARY`` 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. Passthrough is also skipped 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. +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. 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. -Passthrough changes what the rewrite costs, not what it produces, with one -exception worth knowing: a Parquet column chunk can only carry one dictionary, -so when the input files supply different dictionaries the output column keeps -the first and falls back to plain encoding for the rest of the row group. The -rewritten data is unchanged either way, but the output file may be larger than a -rewrite that rebuilt a single dictionary from materialized values. Set -``parquet.read.enable-dictionary-passthrough`` to ``false`` on the table to turn -the optimization off and always rebuild. +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/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index a69b677b9..f819659ad 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -26,7 +26,6 @@ #include "arrow/buffer.h" #include "arrow/c/abi.h" #include "arrow/compute/cast.h" -#include "arrow/compute/exec.h" #include "arrow/type_traits.h" #include "arrow/util/bit_util.h" #include "arrow/util/bitmap_ops.h" @@ -36,6 +35,7 @@ #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 { @@ -43,14 +43,14 @@ 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::IsParquetDictionaryValueType(). +// ArrowUtils::IsDictionaryLayoutRecoverableValueType(). bool IsResolvableDictionary(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::IsParquetDictionaryValueType(*dictionary_type.value_type()); + ArrowUtils::IsDictionaryLayoutRecoverableValueType(*dictionary_type.value_type()); } // Whether `type` is or contains a dictionary at any depth. @@ -552,10 +552,28 @@ Result ArrowUtils::GetCompressionType(const std::strin return compression_type; } -bool ArrowUtils::IsParquetDictionaryValueType(const arrow::DataType& type) { - return type.id() == arrow::Type::STRING || type.id() == arrow::Type::BINARY; +// `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 +// ResolveParquetDictionaryStructType() 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::ResolveParquetDictionaryStructType( const std::shared_ptr& logical_type, const ::ArrowArray* batch) { if (batch == nullptr || logical_type->id() != arrow::Type::STRUCT || @@ -584,7 +602,7 @@ Result> ArrowUtils::ResolveParquetDictionaryStr fields.push_back(field); continue; } - if (!IsParquetDictionaryValueType(*field->type())) { + 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", @@ -607,7 +625,6 @@ Result> ArrowUtils::FlattenUnresolvableDicti return batch; } const auto& logical_struct_type = checked_cast(*logical_type); - arrow::compute::ExecContext exec_context(pool); std::shared_ptr data; arrow::FieldVector fields = batch_type->fields(); for (int32_t i = 0; i < batch_type->num_fields(); ++i) { @@ -629,11 +646,11 @@ Result> ArrowUtils::FlattenUnresolvableDicti } // 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_FROM_ARROW( - arrow::Datum decoded, - arrow::compute::Cast(arrow::MakeArray(data->child_data[i]), logical_field->type(), - arrow::compute::CastOptions::Safe(), &exec_context)); - data->child_data[i] = decoded.array(); + 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) { diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h index e2e6c3ed6..5660cbb4d 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -74,51 +74,28 @@ class PAIMON_EXPORT ArrowUtils { /// Handles "none" and empty string by mapping them to "uncompressed". static Result GetCompressionType(const std::string& compression); - /// Whether a column of `type` may be carried dictionary-encoded across the Arrow C data - /// interface, which drops the type and leaves only the layout behind. + /// 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. /// - /// A layout pins down neither the index width nor the offset width, so the only encoding worth - /// carrying is the one a single known producer emits: `dictionary(int32(), utf8()|binary())`, - /// which is what Arrow's Parquet reader produces for - /// `ArrowReaderProperties::set_read_dictionary`. `LARGE_STRING` is deliberately excluded even - /// though it is binary-like: the ORC reader widens strings to - /// `dictionary(int64(), large_utf8())` under lazy decoding, and 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 - /// ResolveParquetDictionaryStructType() for where the index width becomes a caller contract. - /// - /// This is the single definition shared by the reader that decides which columns to request - /// encoded and by the writer that has to recognise them again on the other side. + /// 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 IsParquetDictionaryValueType(const arrow::DataType& type); + static bool IsDictionaryLayoutRecoverableValueType(const arrow::DataType& type); /// Recovers the struct type of a batch that Arrow's Parquet reader produced with /// `set_read_dictionary` enabled: `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 not inferred, it is assumed, and that assumption is only valid - /// for Arrow's Parquet reader. **The value type check does not make it safe for anything - /// else**: it rejects `dictionary(int64(), large_utf8())`, which is 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 this is a contract, not a check, and it binds the code that *produces* the batch rather - /// than the two places that call this. A producer must either be handing on a batch that came - /// straight from Arrow's Parquet reader, or must run FlattenUnresolvableDictionaries() while - /// the type is still known - that one does test the index width, and decodes every column this - /// cannot resolve while leaving the rest encoded. - /// `AppendOnlyFileStoreWrite::CompactRewrite` is today's only production path that can hand - /// over a batch whose dictionaries the schema does not declare, and it takes the second route. - /// The callers themselves - `ParquetFormatWriter::ResolveBatchSchema` and - /// `DataFileWriterBase::AddFileIndexBatch` - are downstream of it and see only the layout. - /// - /// The value-type rejection and the rejection of a dictionary below the top level narrow the - /// blast radius; they do not close it. Closing it needs the real `ArrowSchema` to reach the - /// writer, which the `FormatWriter::AddBatch(ArrowArray*)` signature currently drops. + /// 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 hand on a + /// batch that came straight from Arrow's Parquet reader, 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. @@ -137,8 +114,7 @@ class PAIMON_EXPORT ArrowUtils { /// resolve stays 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. /// - /// This is the counterpart of the restriction above: exporting an array through the C data - /// interface drops its type, so a column whose encoding does not survive that round trip has + /// The counterpart of the restriction above: an encoding that does not survive the export 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 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 74c964909..78e6c4260 100644 --- a/src/paimon/core/io/data_file_index_writer.cpp +++ b/src/paimon/core/io/data_file_index_writer.cpp @@ -27,7 +27,6 @@ #include "arrow/c/bridge.h" #include "arrow/c/helpers.h" #include "arrow/compute/cast.h" -#include "arrow/compute/exec.h" #include "fmt/format.h" #include "paimon/common/io/byte_array_output_stream.h" #include "paimon/common/io/memory_segment_output_stream.h" @@ -36,6 +35,7 @@ #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" @@ -100,13 +100,12 @@ 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"); } - // Buffers allocated through the adaptor keep a raw pointer to it, so it has to outlive every - // array decoded below. Built on first use, since most batches decode nothing. - std::unique_ptr arrow_pool; // 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 @@ -124,15 +123,13 @@ Status DataFileIndexWriter::AddBatch(const std::shared_ptr& if (cached != decoded_columns.end()) { column = cached->second; } else { - if (arrow_pool == nullptr) { - arrow_pool = GetArrowPool(pool_); + if (arrow_pool_ == nullptr) { + arrow_pool_ = GetArrowPool(pool_); } - arrow::compute::ExecContext exec_context(arrow_pool.get()); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - arrow::Datum decoded, - arrow::compute::Cast(column, entry.field->type(), - arrow::compute::CastOptions::Safe(), &exec_context)); - column = decoded.make_array(); + 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); } } 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 db2306167..9837b5563 100644 --- a/src/paimon/core/io/data_file_index_writer_test.cpp +++ b/src/paimon/core/io/data_file_index_writer_test.cpp @@ -181,21 +181,30 @@ TEST_F(DataFileIndexWriterTest, TestDictionaryEncodedIndexedColumnRoundTrip) { CreateWriter({{"file-index.bitmap.columns", "f0"}, {Options::FILE_INDEX_IN_MANIFEST_THRESHOLD, "1MB"}})); - 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::utf8(), R"(["a", "b", "c"])").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(), "[10, 20, 30, 40]").ValueOrDie(); - auto batch = checked_pointer_cast( - arrow::StructArray::Make({encoded, values}, std::vector{"f0", "f1"}) - .ValueOrDie()); - - ASSERT_OK(writer->AddBatch(batch)); + 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)); @@ -205,9 +214,14 @@ TEST_F(DataFileIndexWriterTest, TestDictionaryEncodedIndexedColumnRoundTrip) { ASSERT_OK_AND_ASSIGN(auto equal_result, bitmap_readers[0]->VisitEqual(Literal(FieldType::STRING, "a", 1))); ASSERT_EQ("{0,2}", equal_result->ToString()); - ASSERT_OK_AND_ASSIGN(auto single_row_result, + // 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}", single_row_result->ToString()); + 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) { 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 7ce2f0a0b..9641a8f6a 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -57,12 +57,10 @@ #include "paimon/core/utils/file_store_path_factory.h" #include "paimon/core/utils/snapshot_manager.h" #include "paimon/executor.h" -#include "paimon/format/parquet/parquet_format_defs.h" #include "paimon/logging.h" #include "paimon/read_context.h" #include "paimon/realtime/realtime_context.h" #include "paimon/result.h" -#include "parquet/properties.h" namespace arrow { class Schema; } // namespace arrow @@ -72,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, @@ -155,10 +169,11 @@ Result>> AppendOnlyFileStoreWrite::Com PAIMON_ASSIGN_OR_RAISE( std::shared_ptr plan_factory, ShreddingWritePlanFactories::SelectActive(options_, write_schema_, pool_)); - PAIMON_ASSIGN_OR_RAISE(bool dictionary_passthrough, CanUseDictionaryPassthrough(plan_factory)); + 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, dictionary_passthrough)); + 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, @@ -299,25 +314,28 @@ Result AppendOnlyFileStoreWrite::GetDat data_file_path_factory, pool_); } -Result AppendOnlyFileStoreWrite::CanUseDictionaryPassthrough( +Result> AppendOnlyFileStoreWrite::GetDictionaryPassthroughVetoReason( const std::shared_ptr& plan_factory) const { std::shared_ptr file_format = options_.GetFileFormat(); - if (!file_format || file_format->Identifier() != "parquet") { - return false; + 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(), parquet::PARQUET_ENABLE_DICTIONARY, - ::parquet::DEFAULT_IS_DICTIONARY_ENABLED)); + PAIMON_ASSIGN_OR_RAISE(bool enable_dictionary, OptionsUtils::GetValueFromMap( + options_.ToMap(), kParquetEnableDictionary, + kDefaultParquetEnableDictionary)); if (!enable_dictionary) { - return false; + return std::optional("parquet.enable-dictionary is false"); } - return plan_factory == nullptr; + 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, bool dictionary_passthrough) const { + const std::vector>& files, + const std::optional& veto_reason) const { ReadContextBuilder context_builder(root_path_); context_builder.SetOptions(options_.ToMap()) .WithFileSystem(options_.GetFileSystem()) @@ -327,16 +345,22 @@ Result> AppendOnlyFileStoreWrite::CreateFilesReader .WithMemoryPool(pool_); PAIMON_ASSIGN_OR_RAISE(std::shared_ptr read_context, context_builder.Finish()); std::map options = options_.ToMap(); - // CompactRewrite copies batches into the rewritten file without looking at any value, so a - // column the input files already store dictionary-encoded can keep that encoding instead of - // being expanded here and hashed again by the writer. - if (dictionary_passthrough) { - // `emplace` so an explicit table option can still turn it off. - options.emplace(parquet::PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH, "true"); - } else { - // Not negotiable the other way: a writer that cannot take a dictionary-encoded batch must - // not receive one because the table happens to set the read option. - options[parquet::PARQUET_READ_ENABLE_DICTIONARY_PASSTHROUGH] = "false"; + // 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)); 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 48dae542b..132118a72 100644 --- a/src/paimon/core/operation/append_only_file_store_write.h +++ b/src/paimon/core/operation/append_only_file_store_write.h @@ -121,7 +121,7 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { /// @param plan_factory The active shredding write plan, or nullptr when the rewrite stays a /// plain passthrough. Resolved by the caller because - /// `CanUseDictionaryPassthrough` needs the same answer. + /// `GetDictionaryPassthroughVetoReason` needs the same answer. Result GetDataFileWriterFactory( const std::shared_ptr& data_file_path_factory, const std::shared_ptr& schema, @@ -129,19 +129,30 @@ class AppendOnlyFileStoreWrite : public AbstractFileStoreWrite { 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, bool dictionary_passthrough) const; + const std::vector>& files, + const std::optional& veto_reason) const; - /// Whether `CompactRewrite` may forward the dictionary encoding of its input files instead of - /// expanding every value. Requires all three of: + /// 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. - Result CanUseDictionaryPassthrough( + /// + /// @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_; diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.cpp b/src/paimon/format/parquet/parquet_file_batch_reader.cpp index aacebadbe..680c0907c 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader.cpp @@ -209,7 +209,18 @@ std::set ParquetFileBatchReader::ResolveFullyDictionaryEncodedColumns( 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); } @@ -268,7 +279,7 @@ std::shared_ptr ParquetFileBatchReader::ApplyDictionaryReadType // 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::IsParquetDictionaryValueType(*field->type())) { + ArrowUtils::IsDictionaryLayoutRecoverableValueType(*field->type())) { fields.push_back(field->WithType(arrow::dictionary(arrow::int32(), field->type()))); } else { fields.push_back(field); diff --git a/src/paimon/format/parquet/parquet_file_batch_reader.h b/src/paimon/format/parquet/parquet_file_batch_reader.h index 95015df87..7c8e3bbdf 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader.h +++ b/src/paimon/format/parquet/parquet_file_batch_reader.h @@ -170,8 +170,9 @@ class ParquetFileBatchReader : public PrefetchFileBatchReader { const std::map& options, int32_t batch_size, const std::optional& hints); - /// Leaf column indices that are candidates for `set_read_dictionary`: non-nested BYTE_ARRAY - /// columns whose every data page, in every row group, is dictionary-encoded. + /// 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); 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 41093eab8..e64fe3650 100644 --- a/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp +++ b/src/paimon/format/parquet/parquet_file_batch_reader_test.cpp @@ -1810,9 +1810,10 @@ TEST_F(ParquetFileBatchReaderTest, TestDictionaryPassthrough) { ASSERT_EQ(arrow::Type::STRING, result.batch->field(1)->type()->id()); } { - // Absent, which is how every read outside the compaction rewrite reaches this reader: the - // default has to be off, or an ordinary scan would start emitting dictionary batches at - // consumers that do not unwrap them. + // 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); @@ -1820,6 +1821,42 @@ TEST_F(ParquetFileBatchReaderTest, TestDictionaryPassthrough) { } } +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 diff --git a/src/paimon/format/parquet/parquet_format_defs.h b/src/paimon/format/parquet/parquet_format_defs.h index dc721d7a2..d22506c45 100644 --- a/src/paimon/format/parquet/parquet_format_defs.h +++ b/src/paimon/format/parquet/parquet_format_defs.h @@ -102,14 +102,15 @@ 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/BINARY columns as Arrow DictionaryArray instead of one copy of -// the value per row. Restricted to non-nested leaf columns whose every data page is already -// dictionary-encoded, so the reader only ever hands on a dictionary the file itself has. +// 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 because it only pays off when the consumer forwards the batch without inspecting -// values, which is why the append compaction rewrite is the one caller that opts in. Value -// accessors have to unwrap DictionaryArray to read such a column; `ColumnarUtils::GetView` does, -// but that is not true of every accessor, so a new consumer has to be checked before enabling it. +// 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"; diff --git a/src/paimon/format/parquet/parquet_format_writer.cpp b/src/paimon/format/parquet/parquet_format_writer.cpp index c2d35d714..05a80f27a 100644 --- a/src/paimon/format/parquet/parquet_format_writer.cpp +++ b/src/paimon/format/parquet/parquet_format_writer.cpp @@ -25,7 +25,7 @@ #include "arrow/array/array_dict.h" #include "arrow/c/bridge.h" -#include "arrow/compute/api.h" +#include "arrow/compute/cast.h" #include "arrow/memory_pool.h" #include "arrow/record_batch.h" #include "arrow/type.h" @@ -36,6 +36,7 @@ #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" @@ -82,25 +83,20 @@ Status ParquetFormatWriter::AddBatch(ArrowArray* batch) { } Result> ParquetFormatWriter::ResolveBatchSchema( - const ::ArrowArray* batch) { + const ::ArrowArray* batch) const { PAIMON_ASSIGN_OR_RAISE( std::shared_ptr batch_type, ArrowUtils::ResolveParquetDictionaryStructType(logical_struct_type_, batch)); if (batch_type == logical_struct_type_) { return schema_; } - if (dictionary_batch_type_ == nullptr || !dictionary_batch_type_->Equals(*batch_type)) { - dictionary_batch_type_ = batch_type; - dictionary_batch_schema_ = arrow::schema(batch_type->fields(), schema_->metadata()); - } - return dictionary_batch_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; - arrow::compute::ExecContext exec_context(pool_.get()); 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 || @@ -112,11 +108,11 @@ Result> ParquetFormatWriter::FlattenUnwritab fields = record_batch->schema()->fields(); } const auto& dictionary_type = checked_cast(*column->type()); - PAIMON_ASSIGN_OR_RAISE_FROM_ARROW( - arrow::Datum flattened, - arrow::compute::Cast(column, dictionary_type.value_type(), - arrow::compute::CastOptions::Safe(), &exec_context)); - columns[i] = flattened.make_array(); + 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()) { diff --git a/src/paimon/format/parquet/parquet_format_writer.h b/src/paimon/format/parquet/parquet_format_writer.h index a0f634397..48f8253de 100644 --- a/src/paimon/format/parquet/parquet_format_writer.h +++ b/src/paimon/format/parquet/parquet_format_writer.h @@ -88,7 +88,7 @@ class ParquetFormatWriter : public FormatWriter { /// 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); + 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 @@ -102,10 +102,6 @@ class ParquetFormatWriter : public FormatWriter { std::shared_ptr schema_; // Struct view of schema_, matched against the layout of each incoming batch. std::shared_ptr logical_struct_type_; - // Last dictionary-encoded batch type and its schema, so a run of identically encoded batches - // builds the import schema only once. - std::shared_ptr dictionary_batch_type_; - std::shared_ptr dictionary_batch_schema_; 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 f18c94d34..3d1c00872 100644 --- a/src/paimon/format/parquet/parquet_format_writer_test.cpp +++ b/src/paimon/format/parquet/parquet_format_writer_test.cpp @@ -856,6 +856,63 @@ TEST_F(ParquetFormatWriterTest, TestGetEstimateLengthWithDictionaryBatches) { 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; diff --git a/test/inte/append_compaction_inte_test.cpp b/test/inte/append_compaction_inte_test.cpp index 0e370e7da..69b12db5e 100644 --- a/test/inte/append_compaction_inte_test.cpp +++ b/test/inte/append_compaction_inte_test.cpp @@ -150,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_; }; @@ -830,9 +858,10 @@ TEST_F(AppendCompactionInteTest, TestAppendTableCompactionWithIOException) { // 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. // -// Parameterised over the two formats that have an encoding to forward or to suppress: Parquet -// turns the passthrough on, ORC forces it off because its writer cannot take a dictionary-encoded -// batch. ORC lazy decoding is on throughout, which makes the ORC reader hand over +// 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)` - a shape no layout can resolve, so it exercises the // decode-at-the-source path rather than the passthrough. TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) { @@ -843,11 +872,13 @@ TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) auto dir = UniqueTestDirectory::Create(); ASSERT_TRUE(dir); - // `s` and `b` are low-cardinality and come back encoded; `id` is INT32, which the gate excludes - // by physical type, so the rewrite carries both kinds of column at once. `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 + // `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()), @@ -859,6 +890,9 @@ TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) {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. @@ -893,13 +927,31 @@ TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) /*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` and `b` arrive as DictionaryArray and - // `id` does not, because the gate only considers BYTE_ARRAY leaves; 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. + // 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)); @@ -907,35 +959,18 @@ TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) auto input_split = std::dynamic_pointer_cast(input_splits[0]); ASSERT_TRUE(input_split); ASSERT_EQ(3, input_split->DataFiles().size()); - std::string input_path = - PathUtil::JoinPath(input_split->BucketPath(), input_split->DataFiles()[0]->file_name); - ASSERT_OK_AND_ASSIGN(auto unique_input_stream, dir->GetFileSystem()->Open(input_path)); - std::shared_ptr input_stream(std::move(unique_input_stream)); - - std::map passthrough_options = options; - passthrough_options["parquet.read.enable-dictionary-passthrough"] = "true"; - ASSERT_OK_AND_ASSIGN(auto input_file_format, - FileFormatFactory::Get(file_format, passthrough_options)); - ASSERT_OK_AND_ASSIGN(auto input_reader_builder, input_file_format->CreateReaderBuilder(10)); - ASSERT_OK_AND_ASSIGN(auto input_reader, input_reader_builder->Build(input_stream)); - ASSERT_OK_AND_ASSIGN(auto c_input_schema, input_reader->GetFileSchema()); - ASSERT_OK(input_reader->SetReadSchema(c_input_schema.get(), /*predicate=*/nullptr, - /*selection_bitmap=*/std::nullopt)); - // Read one batch directly rather than through ReadResultCollector, which decodes - // dictionaries on the way out and would hide the very thing being asserted. - ASSERT_OK_AND_ASSIGN(BatchReader::ReadBatch input_batch, input_reader->NextBatch()); - ASSERT_FALSE(BatchReader::IsEofBatch(input_batch)); - auto& [input_c_array, input_c_schema] = input_batch; - std::shared_ptr input_array = - arrow::ImportArray(input_c_array.get(), input_c_schema.get()).ValueOrDie(); - std::shared_ptr input_type = input_array->type(); + 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_EQ(arrow::Type::DICTIONARY, input_type->field(2)->type()->id()) << "column b"; 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 @@ -944,6 +979,14 @@ TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) *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)); @@ -963,22 +1006,21 @@ TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) ASSERT_TRUE(data_split); ASSERT_EQ(1, data_split->DataFiles().size()); - 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); - ASSERT_OK_AND_ASSIGN(bool success, helper->ReadAndCheckResult(read_type, data_splits, 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"] - ])")); + 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 @@ -1003,20 +1045,24 @@ TEST_P(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthrough) ASSERT_TRUE(expected->Equals(filtered)) << "actual=" << filtered->ToString(); } -// The same rewrite with the passthrough explicitly disabled has to land on the same table, so the -// kill switch is a performance knob and never a correctness one. -TEST_F(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthroughDisabled) { +// 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"}, - {"parquet.read.enable-dictionary-passthrough", "false"}, }; ASSERT_OK_AND_ASSIGN(auto helper, TestHelper::Create(dir->Str(), schema, /*partition_keys=*/{}, /*primary_keys=*/{}, options, @@ -1044,6 +1090,23 @@ TEST_F(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthroughD 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())); @@ -1059,4 +1122,106 @@ TEST_F(AppendCompactionInteTest, TestAppendTableCompactionDictionaryPassthroughD 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 From 02374757184b85ec9376318be474c77d647c5bd8 Mon Sep 17 00:00:00 2001 From: Nicholas Jiang Date: Mon, 31 Aug 2026 20:53:53 +0800 Subject: [PATCH 3/3] perf(parquet): decode every dictionary when the passthrough is vetoed Addresses the review on #257. `FlattenUnresolvableDictionaries()` decided what to decode from whether a dictionary survives the Arrow C data interface, but what decides it is whether the writer receiving the batch recovers an encoding from the layout at all, and only ParquetFormatWriter does. A reader can hand over a dictionary for reasons the Parquet passthrough option does not govern - a format's own lazy-decoding setting - and those stay on when the rewrite vetoes the passthrough. A `dictionary(int32, utf8)` column could therefore reach a writer that imports it against the logical type, where a dictionary child fails on the buffer count. The built-in ORC reader widens to `dictionary(int64, large_utf8)`, which was already decoded for its shape, so reaching this needs a reader outside this repository. CompactRewrite now derives `preserve_layout_recoverable_dictionaries` from the veto: under a veto every dictionary is decoded before the export, whatever its shape. This also stops a dictionary reaching a Parquet writer whose table sets `parquet.enable-dictionary` to false. Rename ArrowUtils::ResolveParquetDictionaryStructType() to ResolveDictionaryStructTypeFromLayout(), and the file-local IsResolvableDictionary() to IsDictionaryLayoutRecoverable(). Neither is Parquet-specific: they read `ArrowArray::dictionary` and the caller's declared type, and DataFileWriterBase::AddFileIndexBatch runs the former for every format. Moving it into the Parquet layer instead would make core depend on Parquet symbols. --- docs/source/user_guide/compaction.rst | 5 + src/paimon/common/utils/arrow/arrow_utils.cpp | 17 ++- src/paimon/common/utils/arrow/arrow_utils.h | 34 +++--- .../common/utils/arrow/arrow_utils_test.cpp | 115 ++++++++++++------ src/paimon/core/io/data_file_writer_base.h | 2 +- .../append_only_file_store_write.cpp | 16 ++- .../format/parquet/parquet_format_writer.cpp | 2 +- test/inte/append_compaction_inte_test.cpp | 5 +- 8 files changed, 133 insertions(+), 63 deletions(-) diff --git a/docs/source/user_guide/compaction.rst b/docs/source/user_guide/compaction.rst index 44419a245..492e27fff 100644 --- a/docs/source/user_guide/compaction.rst +++ b/docs/source/user_guide/compaction.rst @@ -126,6 +126,11 @@ 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. diff --git a/src/paimon/common/utils/arrow/arrow_utils.cpp b/src/paimon/common/utils/arrow/arrow_utils.cpp index f819659ad..45884f5bb 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils.cpp @@ -44,7 +44,7 @@ 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 IsResolvableDictionary(const arrow::DataType& type) { +bool IsDictionaryLayoutRecoverable(const arrow::DataType& type) { if (type.id() != arrow::Type::DICTIONARY) { return false; } @@ -558,7 +558,7 @@ Result ArrowUtils::GetCompressionType(const std::strin // `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 -// ResolveParquetDictionaryStructType() for where the index width becomes a caller contract. +// ResolveDictionaryStructTypeFromLayout() for where the index width becomes a caller contract. bool ArrowUtils::IsDictionaryLayoutRecoverableValueType(const arrow::DataType& type) { return arrow::is_binary_like(type.id()); } @@ -574,7 +574,7 @@ bool ArrowUtils::IsDictionaryLayoutRecoverableValueType(const arrow::DataType& t // 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::ResolveParquetDictionaryStructType( +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()) { @@ -619,7 +619,8 @@ Result> ArrowUtils::ResolveParquetDictionaryStr Result> ArrowUtils::FlattenUnresolvableDictionaries( const std::shared_ptr& batch, - const std::shared_ptr& logical_type, arrow::MemoryPool* pool) { + 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; @@ -629,7 +630,13 @@ Result> ArrowUtils::FlattenUnresolvableDicti arrow::FieldVector fields = batch_type->fields(); for (int32_t i = 0; i < batch_type->num_fields(); ++i) { std::shared_ptr field = fields[i]; - if (IsResolvableDictionary(*field->type()) || !HasDictionary(*field->type())) { + 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 = diff --git a/src/paimon/common/utils/arrow/arrow_utils.h b/src/paimon/common/utils/arrow/arrow_utils.h index 5660cbb4d..f78135073 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -86,14 +86,14 @@ class PAIMON_EXPORT ArrowUtils { /// @return True when `dictionary(int32(), type)` round-trips through an `ArrowArray`. static bool IsDictionaryLayoutRecoverableValueType(const arrow::DataType& type); - /// Recovers the struct type of a batch that Arrow's Parquet reader produced with - /// `set_read_dictionary` enabled: `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. + /// 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 hand on a - /// batch that came straight from Arrow's Parquet reader, or run + /// 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. /// @@ -106,27 +106,33 @@ class PAIMON_EXPORT ArrowUtils { /// @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> ResolveParquetDictionaryStructType( + static Result> ResolveDictionaryStructTypeFromLayout( const std::shared_ptr& logical_type, const ::ArrowArray* batch); - /// Returns `batch` with every top-level column that ResolveParquetDictionaryStructType() could - /// not resolve decoded to the type its field carries in `logical_type`. A column it can - /// resolve stays 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. + /// 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 that does not survive the export has - /// to be decoded while the type is still known. + /// 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); + const std::shared_ptr& logical_type, arrow::MemoryPool* pool, + bool preserve_layout_recoverable_dictionaries); private: static Status InnerCheckNullabilityMatch(const std::shared_ptr& field, diff --git a/src/paimon/common/utils/arrow/arrow_utils_test.cpp b/src/paimon/common/utils/arrow/arrow_utils_test.cpp index bc2266115..bc95a6d68 100644 --- a/src/paimon/common/utils/arrow/arrow_utils_test.cpp +++ b/src/paimon/common/utils/arrow/arrow_utils_test.cpp @@ -927,7 +927,7 @@ TEST(ArrowUtilsTest, TestGetCompressionType) { } } -TEST(ArrowUtilsTest, TestResolveParquetDictionaryStructType) { +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) { @@ -935,7 +935,7 @@ TEST(ArrowUtilsTest, TestResolveParquetDictionaryStructType) { ArrowArrayMarkReleased(&c_array); EXPECT_TRUE(arrow::ExportArray(*array, &c_array).ok()); Result> resolved = - ArrowUtils::ResolveParquetDictionaryStructType(logical_type, &c_array); + ArrowUtils::ResolveDictionaryStructTypeFromLayout(logical_type, &c_array); ArrowArrayRelease(&c_array); return resolved; }; @@ -1046,24 +1046,24 @@ TEST(ArrowUtilsTest, TestResolveParquetDictionaryStructType) { TEST(ArrowUtilsTest, TestFlattenUnresolvableDictionaries) { auto pool = arrow::default_memory_pool(); - auto describable_type = arrow::dictionary(arrow::int32(), arrow::utf8()); + 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 undescribable_type = arrow::dictionary(arrow::int64(), arrow::large_utf8()); + 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 describable = + std::shared_ptr layout_recoverable = arrow::DictionaryArray::FromArrays( - describable_type, + 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 undescribable = + std::shared_ptr layout_unrecoverable = arrow::DictionaryArray::FromArrays( - undescribable_type, + 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()) @@ -1073,12 +1073,13 @@ TEST(ArrowUtilsTest, TestFlattenUnresolvableDictionaries) { // 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({describable, undescribable, ints}, + 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)); + 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()); @@ -1088,31 +1089,34 @@ TEST(ArrowUtilsTest, TestFlattenUnresolvableDictionaries) { .ValueOrDie(); ASSERT_TRUE(flattened->field(1)->Equals(*expected)) << "actual=" << flattened->field(1)->ToString(); - ASSERT_TRUE(flattened->field(0)->Equals(*describable)); + ASSERT_TRUE(flattened->field(0)->Equals(*layout_recoverable)); } { - // A dictionary below the top level is undescribable too, so the whole column is decoded. + // A dictionary below the top level is layout-unrecoverable too, so the whole column is + // decoded. auto nested = - arrow::StructArray::Make({undescribable}, std::vector{"o"}).ValueOrDie(); + 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)); + 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. - // ResolveParquetDictionaryStructType() 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. This is the - // single reason CompactRewrite has to run this before exporting, rather than relying on - // the writer's own rejection. + // 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()), @@ -1121,12 +1125,13 @@ TEST(ArrowUtilsTest, TestFlattenUnresolvableDictionaries) { .ValueOrDie()) .ValueOrDie(); auto batch = checked_pointer_cast( - arrow::StructArray::Make({describable, wide_indices, ints}, + 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)); + 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 = @@ -1141,31 +1146,73 @@ TEST(ArrowUtilsTest, TestFlattenUnresolvableDictionaries) { // 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({describable, describable, ints}, + 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)); + 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({describable, undescribable, ints}, + 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)); + 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/io/data_file_writer_base.h b/src/paimon/core/io/data_file_writer_base.h index bb6b01fd8..9b3d6e71d 100644 --- a/src/paimon/core/io/data_file_writer_base.h +++ b/src/paimon/core/io/data_file_writer_base.h @@ -133,7 +133,7 @@ class DataFileWriterBase : public SingleFileWriter> batch_type = - ArrowUtils::ResolveParquetDictionaryStructType(logical_type_, batch); + 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 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 9641a8f6a..7448a6d96 100644 --- a/src/paimon/core/operation/append_only_file_store_write.cpp +++ b/src/paimon/core/operation/append_only_file_store_write.cpp @@ -219,12 +219,16 @@ Result>> AppendOnlyFileStoreWrite::Com 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 batch layout alone. Decode here, while the type is still known, whatever that - // recovery cannot describe - an ORC reader under lazy decoding hands over - // `dictionary(int64, large_utf8)`, which a layout says nothing about. Only those columns - // pay for it; a Parquet passthrough column stays encoded. - PAIMON_ASSIGN_OR_RAISE(struct_array, ArrowUtils::FlattenUnresolvableDictionaries( - struct_array, logical_type, arrow_pool.get())); + // 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()); diff --git a/src/paimon/format/parquet/parquet_format_writer.cpp b/src/paimon/format/parquet/parquet_format_writer.cpp index 05a80f27a..768615a97 100644 --- a/src/paimon/format/parquet/parquet_format_writer.cpp +++ b/src/paimon/format/parquet/parquet_format_writer.cpp @@ -86,7 +86,7 @@ Result> ParquetFormatWriter::ResolveBatchSchema( const ::ArrowArray* batch) const { PAIMON_ASSIGN_OR_RAISE( std::shared_ptr batch_type, - ArrowUtils::ResolveParquetDictionaryStructType(logical_struct_type_, batch)); + ArrowUtils::ResolveDictionaryStructTypeFromLayout(logical_struct_type_, batch)); if (batch_type == logical_struct_type_) { return schema_; } diff --git a/test/inte/append_compaction_inte_test.cpp b/test/inte/append_compaction_inte_test.cpp index 69b12db5e..4031fa1b0 100644 --- a/test/inte/append_compaction_inte_test.cpp +++ b/test/inte/append_compaction_inte_test.cpp @@ -862,8 +862,9 @@ TEST_F(AppendCompactionInteTest, TestAppendTableCompactionWithIOException) { // 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)` - a shape no layout can resolve, so it exercises the -// decode-at-the-source path rather than the passthrough. +// `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") {