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