From 92d821877ae82c41480befd51c7998bebf98a154 Mon Sep 17 00:00:00 2001 From: Nicholas Jiang Date: Fri, 28 Aug 2026 14:24:47 +0800 Subject: [PATCH] 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 97bb77813..68e8fd48a 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; @@ -484,4 +536,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 d82d84d88..1b34af8a0 100644 --- a/src/paimon/common/utils/arrow/arrow_utils.h +++ b/src/paimon/common/utils/arrow/arrow_utils.h @@ -25,6 +25,8 @@ #include "arrow/util/type_fwd.h" #include "paimon/result.h" +struct ArrowArray; + namespace paimon { class PAIMON_EXPORT ArrowUtils { @@ -69,6 +71,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