diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index 4cde4fb23..e459519aa 100644 --- a/src/iceberg/CMakeLists.txt +++ b/src/iceberg/CMakeLists.txt @@ -326,11 +326,13 @@ if(ICEBERG_BUILD_BUNDLE) avro/avro_schema_util.cc avro/avro_stream_internal.cc parquet/parquet_data_util.cc + parquet/parquet_dictionary_row_group_filter.cc parquet/parquet_metrics.cc parquet/parquet_metrics_row_group_filter.cc parquet/parquet_reader.cc parquet/parquet_register.cc parquet/parquet_schema_util.cc + parquet/parquet_value_util.cc parquet/parquet_writer.cc parquet/row_number_reader.cc) diff --git a/src/iceberg/file_reader.h b/src/iceberg/file_reader.h index 66d40c80c..d6f3b6298 100644 --- a/src/iceberg/file_reader.h +++ b/src/iceberg/file_reader.h @@ -88,7 +88,7 @@ class ICEBERG_EXPORT ReaderProperties : public ConfigBase { /// \brief Validate Arrow batches before merge-on-read projection. /// Default: true. Disable only when the reader is trusted to produce valid arrays. inline static Entry kArrowValidateBatch{"read.arrow.validate-batch", true}; - /// \brief Use footer statistics to prune Parquet row groups. + /// \brief Use footer statistics and dictionaries to prune Parquet row groups. inline static Entry kParquetRowGroupFilter{ "read.parquet.row-group-filter.enabled", true}; /// \brief Skip GenericDatum in Avro reader for better performance. diff --git a/src/iceberg/parquet/parquet_dictionary_row_group_filter.cc b/src/iceberg/parquet/parquet_dictionary_row_group_filter.cc new file mode 100644 index 000000000..21febeab6 --- /dev/null +++ b/src/iceberg/parquet/parquet_dictionary_row_group_filter.cc @@ -0,0 +1,368 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "iceberg/arrow/arrow_status_internal.h" +#include "iceberg/expression/expression_visitor.h" +#include "iceberg/expression/rewrite_not.h" +#include "iceberg/metadata_columns.h" +#include "iceberg/parquet/parquet_dictionary_row_group_filter_internal.h" +#include "iceberg/parquet/parquet_schema_util_internal.h" +#include "iceberg/parquet/parquet_value_util_internal.h" +#include "iceberg/util/macros.h" + +namespace iceberg::parquet { +namespace { + +bool IsDictionaryEncoding(::parquet::Encoding::type encoding) { + return encoding == ::parquet::Encoding::PLAIN_DICTIONARY || + encoding == ::parquet::Encoding::RLE_DICTIONARY; +} + +// A dictionary page alone is insufficient: later pages may have fallen back to +// plain encoding. Follow Java's ParquetUtil.hasNonDictionaryPages rules. +bool IsFullyDictionaryEncoded(const ::parquet::ColumnChunkMetaData& column) { + const auto& stats = column.encoding_stats(); + if (!stats.empty()) { + bool has_dictionary_data = false; + for (const auto& stat : stats) { + if (stat.page_type == ::parquet::PageType::DATA_PAGE || + stat.page_type == ::parquet::PageType::DATA_PAGE_V2) { + if (!IsDictionaryEncoding(stat.encoding)) { + return false; + } + has_dictionary_data |= stat.count > 0; + } + } + return has_dictionary_data; + } + // With legacy encodings, PLAIN_DICTIONARY identifies dictionary data pages. + // With RLE_DICTIONARY, PLAIN could describe either the dictionary or fallback + // data, so encoding statistics are required. + const auto& encodings = column.encodings(); + return std::ranges::find(encodings, ::parquet::Encoding::PLAIN_DICTIONARY) != + encodings.end() && + std::ranges::all_of(encodings, [](auto encoding) { + return encoding == ::parquet::Encoding::PLAIN_DICTIONARY || + encoding == ::parquet::Encoding::RLE || + encoding == ::parquet::Encoding::BIT_PACKED; + }); +} + +template +Result> DecodeDictionary(const ::parquet::DictionaryPage& page, + const ::parquet::ColumnDescriptor& column, + const PrimitiveType& type) { + auto encoding = page.encoding() == ::parquet::Encoding::PLAIN_DICTIONARY + ? ::parquet::Encoding::PLAIN + : page.encoding(); + auto decoder = ::parquet::MakeTypedDecoder(encoding, &column); + decoder->SetData(page.num_values(), page.data(), page.size()); + std::array buffer; + std::vector values; + int remaining = page.num_values(); + while (remaining > 0) { + const int count = std::min(remaining, static_cast(buffer.size())); + if (decoder->Decode(buffer.data(), count) != count) { + return Invalid("Truncated Parquet dictionary page"); + } + for (int i = 0; i < count; ++i) { + auto converted = [&]() { + if constexpr (std::is_same_v) { + return ParquetValueToLiteral(buffer[i], type, column.type_length()); + } else { + return ParquetValueToLiteral(buffer[i], type); + } + }(); + ICEBERG_ASSIGN_OR_RAISE(auto value, std::move(converted)); + values.push_back(std::move(value)); + } + remaining -= count; + } + return values; +} + +class DictionaryVisitor : public BoundVisitor { + public: + DictionaryVisitor(const ::parquet::arrow::SchemaManifest& manifest, + const ::parquet::RowGroupMetaData& row_group, + const std::unordered_map& columns, + const std::shared_ptr<::arrow::io::RandomAccessFile>& input, + const ::parquet::ReaderProperties& properties, + const ::parquet::ApplicationVersion& writer_version) + : manifest_(manifest), + row_group_(row_group), + columns_(columns), + input_(input), + properties_(properties), + writer_version_(writer_version) {} + + Result AlwaysTrue() override { return true; } + Result AlwaysFalse() override { return false; } + // Make rewrites NOT. Negating an inclusive result would be unsafe. + Result Not(bool) override { return true; } + Result And(bool left, bool right) override { return left && right; } + Result Or(bool left, bool right) override { return left || right; } + // Dictionaries contain only non-null values. A rewritten NOT(range) includes + // IS_NULL, which stays inclusive here; footer statistics may rule it out. + Result IsNull(const std::shared_ptr&) override { return true; } + Result NotNull(const std::shared_ptr&) override { return true; } + + Result IsNaN(const std::shared_ptr& expr) override { + return AnyMatch(expr, [](const Literal& value) { return value.IsNaN(); }); + } + Result NotNaN(const std::shared_ptr& expr) override { + return AnyMatch(expr, [](const Literal& value) { return !value.IsNaN(); }, true); + } + + Result Lt(const std::shared_ptr& expr, const Literal& value) override { + return AnyMatch(expr, [&](const Literal& item) { + return (item <=> value) == std::partial_ordering::unordered || item < value; + }); + } + + Result LtEq(const std::shared_ptr& expr, const Literal& value) override { + return AnyMatch(expr, [&](const Literal& item) { + return (item <=> value) == std::partial_ordering::unordered || item <= value; + }); + } + + Result Gt(const std::shared_ptr& expr, const Literal& value) override { + return AnyMatch(expr, [&](const Literal& item) { + return (item <=> value) == std::partial_ordering::unordered || item > value; + }); + } + + Result GtEq(const std::shared_ptr& expr, const Literal& value) override { + return AnyMatch(expr, [&](const Literal& item) { + return (item <=> value) == std::partial_ordering::unordered || item >= value; + }); + } + + Result Eq(const std::shared_ptr& expr, const Literal& value) override { + return AnyMatch(expr, [&](const Literal& item) { + return (item <=> value) == std::partial_ordering::unordered || item == value; + }); + } + + Result NotEq(const std::shared_ptr& expr, const Literal& value) override { + return AnyMatch( + expr, + [&](const Literal& item) { + return (item <=> value) == std::partial_ordering::unordered || item != value; + }, + true); + } + + Result In(const std::shared_ptr& expr, + const BoundSetPredicate::LiteralSet& values) override { + return AnyMatch(expr, [&](const Literal& item) { return values.contains(item); }); + } + + Result NotIn(const std::shared_ptr& expr, + const BoundSetPredicate::LiteralSet& values) override { + return AnyMatch( + expr, [&](const Literal& item) { return !values.contains(item); }, true); + } + + Result StartsWith(const std::shared_ptr& expr, + const Literal& value) override { + return AnyMatch(expr, [&](const Literal& item) { + return std::get(item.value()) + .starts_with(std::get(value.value())); + }); + } + + Result NotStartsWith(const std::shared_ptr& expr, + const Literal& value) override { + return AnyMatch( + expr, + [&](const Literal& item) { + return !std::get(item.value()) + .starts_with(std::get(value.value())); + }, + true); + } + + private: + template + Result AnyMatch(const std::shared_ptr& expr, Predicate predicate, + bool nulls_match = false) { + auto ref = std::dynamic_pointer_cast(expr); + ICEBERG_ASSIGN_OR_RAISE(auto dictionary, GetDictionary(ref)); + if (!dictionary || (nulls_match && MayContainNull(*ref))) { + return true; + } + return std::ranges::any_of(*dictionary, predicate); + } + + bool MayContainNull(const BoundReference& ref) const { + const auto column = row_group_.ColumnChunk(columns_.at(ref.field_id())); + try { + auto stats = column->statistics(); + return !stats || !stats->HasNullCount() || stats->null_count() != 0; + } catch (const ::parquet::ParquetException&) { + // Optional, unusable statistics cannot prove that nulls are absent. + return true; + } + } + + Result*> GetDictionary( + const std::shared_ptr& ref) { + if (!ref || !ref->type()->is_primitive() || + ref->type()->type_id() == TypeId::kUnknown || + MetadataColumns::IsMetadataColumn(ref->field_id())) { + return nullptr; + } + auto index = columns_.find(ref->field_id()); + if (index == columns_.end()) { + // Missing fields may have initial defaults. + return nullptr; + } + auto field = manifest_.column_index_to_field.find(index->second); + if (field == manifest_.column_index_to_field.end() || + !ValidateParquetTypeCompatibility(*ref->type(), *field->second)) { + return nullptr; + } + auto [cached, inserted] = dictionaries_.try_emplace(ref->field_id()); + if (inserted) { + const auto column = row_group_.ColumnChunk(index->second); + const auto* descriptor = manifest_.descr->Column(index->second); + // A raw dictionary slice has no Parquet column decryption context. Leave + // encrypted columns to the normal reader rather than interpreting ciphertext. + if (descriptor->max_repetition_level() > 0 || + descriptor->physical_type() == ::parquet::Type::INT96 || + descriptor->physical_type() == ::parquet::Type::BOOLEAN || + column->crypto_metadata() || !IsFullyDictionaryEncoded(*column)) { + return nullptr; + } + // GetColumnPageReader eagerly reads the entire column chunk by default. + // Restrict planning I/O to the dictionary page, including its header. + const int64_t start = column->dictionary_page_offset(); + const int64_t end = column->data_page_offset(); + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto file_size, input_->GetSize()); + if (!column->has_dictionary_page() || start < 4 || end <= start || + end > file_size) { + return Invalid("Invalid Parquet dictionary offsets [{}, {}) for field {}", start, + end, ref->field_id()); + } + ICEBERG_ARROW_ASSIGN_OR_RETURN(auto data, input_->ReadAt(start, end - start)); + if (data->size() != end - start) { + return Invalid("Truncated Parquet dictionary page for field {}", ref->field_id()); + } + auto stream = std::make_shared<::arrow::io::BufferReader>(std::move(data)); + // Preserve Arrow's compatibility behavior for old compressed page headers. + bool always_compressed = writer_version_.VersionLt( + ::parquet::ApplicationVersion::PARQUET_CPP_10353_FIXED_VERSION()); + auto pages = + ::parquet::PageReader::Open(stream, column->num_values(), column->compression(), + properties_, *descriptor, always_compressed); + auto page = pages->NextPage(); + if (!page || page->type() != ::parquet::PageType::DICTIONARY_PAGE) { + return Invalid("Missing Parquet dictionary page for field {}", ref->field_id()); + } + const auto& dictionary = static_cast(*page); + if (dictionary.num_values() < 0) { + return Invalid("Negative Parquet dictionary size for field {}", ref->field_id()); + } + const auto& type = static_cast(*ref->type()); + auto decoded = [&]() -> Result> { + switch (descriptor->physical_type()) { + case ::parquet::Type::INT32: + return DecodeDictionary<::parquet::Int32Type>(dictionary, *descriptor, type); + case ::parquet::Type::INT64: + return DecodeDictionary<::parquet::Int64Type>(dictionary, *descriptor, type); + case ::parquet::Type::FLOAT: + return DecodeDictionary<::parquet::FloatType>(dictionary, *descriptor, type); + case ::parquet::Type::DOUBLE: + return DecodeDictionary<::parquet::DoubleType>(dictionary, *descriptor, type); + case ::parquet::Type::BYTE_ARRAY: + return DecodeDictionary<::parquet::ByteArrayType>(dictionary, *descriptor, + type); + case ::parquet::Type::FIXED_LEN_BYTE_ARRAY: + return DecodeDictionary<::parquet::FLBAType>(dictionary, *descriptor, type); + default: + return NotSupported("Unsupported Parquet dictionary type"); + } + }(); + ICEBERG_ASSIGN_OR_RAISE(auto values, std::move(decoded)); + cached->second = std::move(values); + } + return cached->second ? &*cached->second : nullptr; + } + + const ::parquet::arrow::SchemaManifest& manifest_; + const ::parquet::RowGroupMetaData& row_group_; + const std::unordered_map& columns_; + const std::shared_ptr<::arrow::io::RandomAccessFile>& input_; + const ::parquet::ReaderProperties& properties_; + const ::parquet::ApplicationVersion& writer_version_; + std::unordered_map>> dictionaries_; +}; + +} // namespace + +Result> +ParquetDictionaryRowGroupFilter::Make(const std::shared_ptr& filter, + const ::parquet::SchemaDescriptor& file_schema) { + auto result = std::unique_ptr( + new ParquetDictionaryRowGroupFilter()); + for (int i = 0; i < file_schema.num_columns(); ++i) { + auto id = file_schema.Column(i)->schema_node()->field_id(); + if (id >= 0) { + result->field_id_to_column_idx_.emplace(id, i); + } + } + result->bound_ = True::Instance(); + if (filter) { + ICEBERG_ASSIGN_OR_RAISE(result->bound_, RewriteNot::Visit(filter)); + } + return result; +} + +Result ParquetDictionaryRowGroupFilter::ShouldRead( + const ::parquet::arrow::SchemaManifest& manifest, + const ::parquet::FileMetaData& metadata, int row_group_index, + const std::shared_ptr<::arrow::io::RandomAccessFile>& input, + const ::parquet::ReaderProperties& properties) const { + auto row_group = metadata.RowGroup(row_group_index); + if (row_group->num_rows() <= 0) { + return false; + } + try { + DictionaryVisitor visitor(manifest, *row_group, field_id_to_column_idx_, input, + properties, metadata.writer_version()); + return Visit(bound_, visitor); + } catch (const ::parquet::ParquetException& e) { + return Invalid("Cannot read Parquet dictionary: {}", e.what()); + } +} + +} // namespace iceberg::parquet diff --git a/src/iceberg/parquet/parquet_dictionary_row_group_filter_internal.h b/src/iceberg/parquet/parquet_dictionary_row_group_filter_internal.h new file mode 100644 index 000000000..08ca355cf --- /dev/null +++ b/src/iceberg/parquet/parquet_dictionary_row_group_filter_internal.h @@ -0,0 +1,56 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include +#include +#include + +#include +#include + +#include "iceberg/iceberg_bundle_export.h" +#include "iceberg/result.h" +#include "iceberg/type_fwd.h" + +namespace iceberg::parquet { + +// Inclusive dictionary filtering, following Java's +// ParquetDictionaryRowGroupFilter. False proves that the row group cannot match; +// true still requires evaluating the residual predicate on the returned rows. +class ICEBERG_BUNDLE_EXPORT ParquetDictionaryRowGroupFilter { + public: + static Result> Make( + const std::shared_ptr& filter, + const ::parquet::SchemaDescriptor& file_schema); + + // Reads only dictionary-page byte ranges. Input and properties must match metadata. + Result ShouldRead(const ::parquet::arrow::SchemaManifest& manifest, + const ::parquet::FileMetaData& metadata, int row_group_index, + const std::shared_ptr<::arrow::io::RandomAccessFile>& input, + const ::parquet::ReaderProperties& properties) const; + + private: + ParquetDictionaryRowGroupFilter() = default; + std::shared_ptr bound_; + std::unordered_map field_id_to_column_idx_; +}; + +} // namespace iceberg::parquet diff --git a/src/iceberg/parquet/parquet_metrics.cc b/src/iceberg/parquet/parquet_metrics.cc index bad4545a7..f64ccfa2f 100644 --- a/src/iceberg/parquet/parquet_metrics.cc +++ b/src/iceberg/parquet/parquet_metrics.cc @@ -34,13 +34,12 @@ #include "iceberg/expression/literal.h" #include "iceberg/parquet/parquet_metrics_internal.h" +#include "iceberg/parquet/parquet_value_util_internal.h" #include "iceberg/result.h" #include "iceberg/schema.h" #include "iceberg/type.h" #include "iceberg/util/checked_cast.h" -#include "iceberg/util/decimal.h" #include "iceberg/util/truncate_util.h" -#include "iceberg/util/uuid.h" #include "iceberg/util/visit_type.h" namespace iceberg::parquet { @@ -86,99 +85,6 @@ Result TypedStatsLiteral(const ::parquet::Statistics& stats, bool is_mi return converter(is_min ? typed_stats.min() : typed_stats.max()); } -std::vector BytesFromByteArray(const ::parquet::ByteArray& value) { - return std::vector{value.ptr, value.ptr + value.len}; -} - -std::vector BytesFromFLBA(const ::parquet::FixedLenByteArray& value, - int32_t length) { - return std::vector{value.ptr, value.ptr + length}; -} - -Literal DecimalLiteral(int128_t value, const PrimitiveType& iceberg_type) { - const auto& decimal_type = internal::checked_cast(iceberg_type); - return Literal::Decimal(value, decimal_type.precision(), decimal_type.scale()); -} - -Result DecimalLiteralFromBytes(std::span bytes, - const PrimitiveType& iceberg_type) { - ICEBERG_ASSIGN_OR_RAISE(auto decimal, - Decimal::FromBigEndian(bytes.data(), bytes.size())); - return DecimalLiteral(decimal.value(), iceberg_type); -} - -Result BinaryStatsLiteral(std::vector bytes, - const PrimitiveType& iceberg_type) { - switch (iceberg_type.type_id()) { - case TypeId::kString: - return Literal::String(std::string(bytes.begin(), bytes.end())); - case TypeId::kBinary: - return Literal::Binary(std::move(bytes)); - case TypeId::kFixed: - return Literal::Fixed(std::move(bytes)); - case TypeId::kUuid: { - ICEBERG_ASSIGN_OR_RAISE(auto uuid, Uuid::FromBytes(bytes)); - return Literal::UUID(std::move(uuid)); - } - case TypeId::kDecimal: - return DecimalLiteralFromBytes(bytes, iceberg_type); - default: - return InvalidArgument( - "Cannot convert Parquet binary statistics to Iceberg type {}", - iceberg_type.ToString()); - } -} - -Result Int32StatsLiteral(int32_t value, const PrimitiveType& iceberg_type) { - switch (iceberg_type.type_id()) { - case TypeId::kInt: - return Literal::Int(value); - case TypeId::kLong: - return Literal::Long(value); - case TypeId::kDate: - return Literal::Date(value); - case TypeId::kDecimal: - return DecimalLiteral(value, iceberg_type); - default: - return InvalidArgument("Cannot convert Parquet INT32 statistics to Iceberg type {}", - iceberg_type.ToString()); - } -} - -Result Int64StatsLiteral(int64_t value, const PrimitiveType& iceberg_type) { - switch (iceberg_type.type_id()) { - case TypeId::kLong: - return Literal::Long(value); - case TypeId::kTime: - return Literal::Time(value); - case TypeId::kTimestamp: - return Literal::Timestamp(value); - case TypeId::kTimestampTz: - return Literal::TimestampTz(value); - case TypeId::kTimestampNs: - return Literal::TimestampNs(value); - case TypeId::kTimestampTzNs: - return Literal::TimestampTzNs(value); - case TypeId::kDecimal: - return DecimalLiteral(value, iceberg_type); - default: - return InvalidArgument("Cannot convert Parquet INT64 statistics to Iceberg type {}", - iceberg_type.ToString()); - } -} - -Result FloatStatsLiteral(float value, const PrimitiveType& iceberg_type) { - switch (iceberg_type.type_id()) { - case TypeId::kFloat: - return Literal::Float(value); - case TypeId::kDouble: - return Literal::Double(value); - default: - return InvalidArgument("Cannot convert Parquet FLOAT statistics to Iceberg type {}", - iceberg_type.ToString()); - } -} - bool IsFloatingType(const PrimitiveType& type) { return type.type_id() == TypeId::kFloat || type.type_id() == TypeId::kDouble; } @@ -468,33 +374,34 @@ Result ParquetMetrics::StatsValueToLiteral( const ::parquet::Statistics& stats, bool is_min) { switch (column.physical_type()) { case ::parquet::Type::BOOLEAN: - return TypedStatsLiteral<::parquet::BoolStatistics>( - stats, is_min, [](bool value) { return Literal::Boolean(value); }); + return TypedStatsLiteral<::parquet::BoolStatistics>(stats, is_min, [&](bool value) { + return ParquetValueToLiteral(value, iceberg_type); + }); case ::parquet::Type::INT32: return TypedStatsLiteral<::parquet::Int32Statistics>( stats, is_min, - [&](int32_t value) { return Int32StatsLiteral(value, iceberg_type); }); + [&](int32_t value) { return ParquetValueToLiteral(value, iceberg_type); }); case ::parquet::Type::INT64: return TypedStatsLiteral<::parquet::Int64Statistics>( stats, is_min, - [&](int64_t value) { return Int64StatsLiteral(value, iceberg_type); }); + [&](int64_t value) { return ParquetValueToLiteral(value, iceberg_type); }); case ::parquet::Type::FLOAT: return TypedStatsLiteral<::parquet::FloatStatistics>( stats, is_min, - [&](float value) { return FloatStatsLiteral(value, iceberg_type); }); + [&](float value) { return ParquetValueToLiteral(value, iceberg_type); }); case ::parquet::Type::DOUBLE: return TypedStatsLiteral<::parquet::DoubleStatistics>( - stats, is_min, [](double value) { return Literal::Double(value); }); + stats, is_min, + [&](double value) { return ParquetValueToLiteral(value, iceberg_type); }); case ::parquet::Type::BYTE_ARRAY: return TypedStatsLiteral<::parquet::ByteArrayStatistics>( stats, is_min, [&](const ::parquet::ByteArray& value) { - return BinaryStatsLiteral(BytesFromByteArray(value), iceberg_type); + return ParquetValueToLiteral(value, iceberg_type); }); case ::parquet::Type::FIXED_LEN_BYTE_ARRAY: return TypedStatsLiteral<::parquet::FLBAStatistics>( stats, is_min, [&](const ::parquet::FixedLenByteArray& value) { - return BinaryStatsLiteral(BytesFromFLBA(value, column.type_length()), - iceberg_type); + return ParquetValueToLiteral(value, iceberg_type, column.type_length()); }); case ::parquet::Type::INT96: case ::parquet::Type::UNDEFINED: diff --git a/src/iceberg/parquet/parquet_reader.cc b/src/iceberg/parquet/parquet_reader.cc index 47003305a..a59947ab0 100644 --- a/src/iceberg/parquet/parquet_reader.cc +++ b/src/iceberg/parquet/parquet_reader.cc @@ -42,6 +42,7 @@ #include "iceberg/arrow/metadata_column_util_internal.h" #include "iceberg/expression/binder.h" #include "iceberg/parquet/parquet_data_util_internal.h" +#include "iceberg/parquet/parquet_dictionary_row_group_filter_internal.h" #include "iceberg/parquet/parquet_metrics_row_group_filter_internal.h" #include "iceberg/parquet/parquet_register.h" #include "iceberg/parquet/parquet_schema_util_internal.h" @@ -133,6 +134,7 @@ class ParquetReader::Impl { read_schema_ = options.projection; stats_filter_.reset(); + dictionary_filter_.reset(); // Prepare reader properties ::parquet::ReaderProperties reader_properties(pool_); @@ -165,9 +167,12 @@ class ParquetReader::Impl { } ICEBERG_ASSIGN_OR_RAISE(stats_filter_, ParquetMetricsRowGroupFilter::Make( filter, *reader_->manifest().descr)); + ICEBERG_ASSIGN_OR_RAISE( + dictionary_filter_, + ParquetDictionaryRowGroupFilter::Make(filter, *reader_->manifest().descr)); } - ICEBERG_RETURN_UNEXPECTED(PlanRowGroups(options)); + ICEBERG_RETURN_UNEXPECTED(PlanRowGroups(options, reader_properties)); metadata_context_ = {.file_path = options.path, .first_row_id = options.first_row_id, @@ -257,7 +262,8 @@ class ParquetReader::Impl { } private: - Status PlanRowGroups(const ReaderOptions& options) { + Status PlanRowGroups(const ReaderOptions& options, + const ::parquet::ReaderProperties& reader_properties) { auto metadata = reader_->parquet_reader()->metadata(); if (options.split.has_value()) { const auto& split = *options.split; @@ -300,6 +306,12 @@ class ParquetReader::Impl { auto row_group = metadata->RowGroup(row_group_index); ICEBERG_ASSIGN_OR_RAISE( auto should_read, stats_filter_->ShouldRead(reader_->manifest(), *row_group)); + if (should_read && dictionary_filter_) { + ICEBERG_ASSIGN_OR_RAISE( + should_read, dictionary_filter_->ShouldRead(reader_->manifest(), *metadata, + row_group_index, input_stream_, + reader_properties)); + } if (should_read) { row_group_indices_[selected_count++] = row_group_index; } @@ -357,6 +369,7 @@ class ParquetReader::Impl { std::shared_ptr<::iceberg::Schema> read_schema_; // Footer-statistics filter for row group pruning, if enabled. std::unique_ptr stats_filter_; + std::unique_ptr dictionary_filter_; // The projection result to apply to the read schema. SchemaProjection projection_; // The input stream to read Parquet file. diff --git a/src/iceberg/parquet/parquet_value_util.cc b/src/iceberg/parquet/parquet_value_util.cc new file mode 100644 index 000000000..d4739548c --- /dev/null +++ b/src/iceberg/parquet/parquet_value_util.cc @@ -0,0 +1,141 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#include "iceberg/parquet/parquet_value_util_internal.h" +#include "iceberg/type.h" +#include "iceberg/util/checked_cast.h" +#include "iceberg/util/decimal.h" +#include "iceberg/util/uuid.h" + +namespace iceberg::parquet { +namespace { + +std::vector BytesFromByteArray(const ::parquet::ByteArray& value) { + return std::vector{value.ptr, value.ptr + value.len}; +} + +std::vector BytesFromFLBA(const ::parquet::FixedLenByteArray& value, + int32_t length) { + return std::vector{value.ptr, value.ptr + length}; +} + +Literal DecimalLiteral(int128_t value, const PrimitiveType& iceberg_type) { + const auto& decimal_type = internal::checked_cast(iceberg_type); + return Literal::Decimal(value, decimal_type.precision(), decimal_type.scale()); +} + +Result DecimalLiteralFromBytes(std::span bytes, + const PrimitiveType& iceberg_type) { + ICEBERG_ASSIGN_OR_RAISE(auto decimal, + Decimal::FromBigEndian(bytes.data(), bytes.size())); + return DecimalLiteral(decimal.value(), iceberg_type); +} + +Result BinaryLiteral(std::vector bytes, + const PrimitiveType& iceberg_type) { + switch (iceberg_type.type_id()) { + case TypeId::kString: + return Literal::String(std::string(bytes.begin(), bytes.end())); + case TypeId::kBinary: + return Literal::Binary(std::move(bytes)); + case TypeId::kFixed: + return Literal::Fixed(std::move(bytes)); + case TypeId::kUuid: { + ICEBERG_ASSIGN_OR_RAISE(auto uuid, Uuid::FromBytes(bytes)); + return Literal::UUID(std::move(uuid)); + } + case TypeId::kDecimal: + return DecimalLiteralFromBytes(bytes, iceberg_type); + default: + return InvalidArgument("Cannot convert Parquet binary value to Iceberg type {}", + iceberg_type.ToString()); + } +} + +} // namespace + +Result ParquetValueToLiteral(int32_t value, const PrimitiveType& iceberg_type) { + switch (iceberg_type.type_id()) { + case TypeId::kInt: + return Literal::Int(value); + case TypeId::kLong: + return Literal::Long(value); + case TypeId::kDate: + return Literal::Date(value); + case TypeId::kDecimal: + return DecimalLiteral(value, iceberg_type); + default: + return InvalidArgument("Cannot convert Parquet INT32 value to Iceberg type {}", + iceberg_type.ToString()); + } +} + +Result ParquetValueToLiteral(int64_t value, const PrimitiveType& iceberg_type) { + switch (iceberg_type.type_id()) { + case TypeId::kLong: + return Literal::Long(value); + case TypeId::kTime: + return Literal::Time(value); + case TypeId::kTimestamp: + return Literal::Timestamp(value); + case TypeId::kTimestampTz: + return Literal::TimestampTz(value); + case TypeId::kTimestampNs: + return Literal::TimestampNs(value); + case TypeId::kTimestampTzNs: + return Literal::TimestampTzNs(value); + case TypeId::kDecimal: + return DecimalLiteral(value, iceberg_type); + default: + return InvalidArgument("Cannot convert Parquet INT64 value to Iceberg type {}", + iceberg_type.ToString()); + } +} + +Result ParquetValueToLiteral(float value, const PrimitiveType& iceberg_type) { + switch (iceberg_type.type_id()) { + case TypeId::kFloat: + return Literal::Float(value); + case TypeId::kDouble: + return Literal::Double(value); + default: + return InvalidArgument("Cannot convert Parquet FLOAT value to Iceberg type {}", + iceberg_type.ToString()); + } +} + +Result ParquetValueToLiteral(bool value, const PrimitiveType&) { + return Literal::Boolean(value); +} + +Result ParquetValueToLiteral(double value, const PrimitiveType&) { + return Literal::Double(value); +} + +Result ParquetValueToLiteral(const ::parquet::ByteArray& value, + const PrimitiveType& iceberg_type) { + return BinaryLiteral(BytesFromByteArray(value), iceberg_type); +} + +Result ParquetValueToLiteral(const ::parquet::FixedLenByteArray& value, + const PrimitiveType& iceberg_type, int32_t length) { + return BinaryLiteral(BytesFromFLBA(value, length), iceberg_type); +} + +} // namespace iceberg::parquet diff --git a/src/iceberg/parquet/parquet_value_util_internal.h b/src/iceberg/parquet/parquet_value_util_internal.h new file mode 100644 index 000000000..8d03a871d --- /dev/null +++ b/src/iceberg/parquet/parquet_value_util_internal.h @@ -0,0 +1,41 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +#pragma once + +#include + +#include "iceberg/expression/literal.h" +#include "iceberg/result.h" + +namespace iceberg::parquet { + +// Convert physical Parquet values to owned Iceberg literals after validating the +// file type against the expected Iceberg type. Shared by statistics and dictionaries. +Result ParquetValueToLiteral(bool value, const PrimitiveType& iceberg_type); +Result ParquetValueToLiteral(int32_t value, const PrimitiveType& iceberg_type); +Result ParquetValueToLiteral(int64_t value, const PrimitiveType& iceberg_type); +Result ParquetValueToLiteral(float value, const PrimitiveType& iceberg_type); +Result ParquetValueToLiteral(double value, const PrimitiveType& iceberg_type); +Result ParquetValueToLiteral(const ::parquet::ByteArray& value, + const PrimitiveType& iceberg_type); +Result ParquetValueToLiteral(const ::parquet::FixedLenByteArray& value, + const PrimitiveType& iceberg_type, int32_t length); + +} // namespace iceberg::parquet diff --git a/src/iceberg/test/parquet_row_group_filter_test.cc b/src/iceberg/test/parquet_row_group_filter_test.cc index cf44af16c..580176865 100644 --- a/src/iceberg/test/parquet_row_group_filter_test.cc +++ b/src/iceberg/test/parquet_row_group_filter_test.cc @@ -17,15 +17,22 @@ * under the License. */ +#include +#include +#include + #include #include #include #include +#include #include #include #include +#include #include #include +#include #include #include @@ -37,6 +44,7 @@ #include "iceberg/metadata_columns.h" #include "iceberg/name_mapping.h" #include "iceberg/parquet/parquet_data_util_internal.h" +#include "iceberg/parquet/parquet_dictionary_row_group_filter_internal.h" #include "iceberg/parquet/parquet_metrics_row_group_filter_internal.h" #include "iceberg/parquet/parquet_register.h" #include "iceberg/schema.h" @@ -49,6 +57,66 @@ namespace iceberg::parquet { namespace { +struct DictionaryReadRange { + int64_t offset; + int64_t length; + bool operator==(const DictionaryReadRange&) const = default; +}; + +class DictionaryRecordingStream : public SeekableInputStream { + public: + DictionaryRecordingStream(std::unique_ptr stream, + std::vector& reads) + : stream_(std::move(stream)), reads_(reads) {} + Result Position() const override { return stream_->Position(); } + Status Seek(int64_t position) override { return stream_->Seek(position); } + Result Read(std::span out) override { + ICEBERG_ASSIGN_OR_RAISE(auto position, stream_->Position()); + reads_.push_back({.offset = position, .length = static_cast(out.size())}); + return stream_->Read(out); + } + Status ReadFully(int64_t position, std::span out) override { + reads_.push_back({.offset = position, .length = static_cast(out.size())}); + return stream_->ReadFully(position, out); + } + Status Close() override { return stream_->Close(); } + + private: + std::unique_ptr stream_; + std::vector& reads_; +}; + +class DictionaryRecordingInput : public InputFile { + public: + DictionaryRecordingInput(std::unique_ptr input, + std::vector& reads) + : input_(std::move(input)), reads_(reads) {} + std::string_view location() const override { return input_->location(); } + Result Size() const override { return input_->Size(); } + Result> Open() override { + ICEBERG_ASSIGN_OR_RAISE(auto stream, input_->Open()); + return std::make_unique(std::move(stream), reads_); + } + + private: + std::unique_ptr input_; + std::vector& reads_; +}; + +class DictionaryRecordingFileIO : public MockFileIO { + public: + Result> NewInputFile(std::string path) override { + ICEBERG_ASSIGN_OR_RAISE(auto file, MockFileIO::NewInputFile(std::move(path))); + return std::make_unique(std::move(file), reads); + } + Result> NewInputFile(std::string path, + size_t length) override { + ICEBERG_ASSIGN_OR_RAISE(auto file, MockFileIO::NewInputFile(std::move(path), length)); + return std::make_unique(std::move(file), reads); + } + std::vector reads; +}; + class ParquetRowGroupFilterTest : public ::testing::Test { protected: void SetUp() override { @@ -85,11 +153,15 @@ class ParquetRowGroupFilterTest : public ::testing::Test { ICEBERG_ASSIGN_OR_RAISE(auto out, arrow::OpenArrowOutputStream(io_, path_)); ::parquet::WriterProperties::Builder properties; properties.disable_dictionary(); + if (writer_properties_) { + properties = ::parquet::WriterProperties::Builder(*writer_properties_); + } if (!statistics) { properties.disable_statistics(); } - ICEBERG_ARROW_RETURN_NOT_OK(::parquet::arrow::WriteTable( - *table, ::arrow::default_memory_pool(), out, 2, properties.build())); + ICEBERG_ARROW_RETURN_NOT_OK( + ::parquet::arrow::WriteTable(*table, ::arrow::default_memory_pool(), out, + row_group_size_, properties.build())); ICEBERG_ARROW_RETURN_NOT_OK(out->Close()); ICEBERG_ASSIGN_OR_RAISE(auto input, arrow::OpenArrowInputStream(io_, path_)); metadata_ = ::parquet::ParquetFileReader::Open(input)->metadata(); @@ -153,6 +225,8 @@ class ParquetRowGroupFilterTest : public ::testing::Test { EXPECT_EQ(actual, expected); } + int64_t row_group_size_ = 2; + std::shared_ptr<::parquet::WriterProperties> writer_properties_; std::string path_ = "rg.parquet"; std::shared_ptr io_; std::shared_ptr schema_; @@ -855,5 +929,652 @@ TEST_F(ParquetRowGroupFilterTest, DoubleBoundsMustNotHideNaNs) { Check(Options(Expressions::LessThan("key", Literal::Double(-3))), {0, 1, 2, 3}); } +// Exercise dictionary filtering separately from footer statistics, and through +// the production reader (where the filter column is outside the projection). +class ParquetDictionaryFilterTest : public ParquetRowGroupFilterTest { + protected: + void SetUp() override { + ParquetRowGroupFilterTest::SetUp(); + io_ = std::make_shared(); + ::parquet::WriterProperties::Builder properties; + writer_properties_ = properties.enable_dictionary()->build(); + } + + void CheckDictionaries(const std::shared_ptr& expression, + const std::vector& expected) { + ICEBERG_UNWRAP_OR_FAIL(auto input, arrow::OpenArrowInputStream(io_, path_)); + ::parquet::arrow::SchemaManifest manifest; + ::parquet::ArrowReaderProperties properties; + properties.set_arrow_extensions_enabled(true); + ASSERT_TRUE(::parquet::arrow::SchemaManifest::Make(metadata_->schema(), + metadata_->key_value_metadata(), + properties, &manifest) + .ok()); + ICEBERG_UNWRAP_OR_FAIL( + auto filter, ParquetDictionaryRowGroupFilter::Make(Options(expression).filter, + *metadata_->schema())); + ASSERT_EQ(expected.size(), metadata_->num_row_groups()); + for (int i = 0; i < metadata_->num_row_groups(); ++i) { + SCOPED_TRACE(expression->ToString() + " group " + std::to_string(i)); + ICEBERG_UNWRAP_OR_FAIL(auto actual, + filter->ShouldRead(manifest, *metadata_, i, input, + ::parquet::default_reader_properties())); + EXPECT_EQ(actual, expected[i]); + } + } + + // Rebuild the footer while retaining the real file's pages and encoding stats. + std::unique_ptr<::parquet::FileMetaData> RebuildMetadata( + bool corrupt_stats = false, bool encoding_stats = true, + std::optional dictionary_offset = std::nullopt, + std::optional data_offset = std::nullopt) { + auto builder = + ::parquet::FileMetaDataBuilder::Make(metadata_->schema(), writer_properties_); + for (int rg = 0; rg < metadata_->num_row_groups(); ++rg) { + auto original = metadata_->RowGroup(rg); + auto group = builder->AppendRowGroup(); + group->set_num_rows(original->num_rows()); + for (int i = 0; i < original->num_columns(); ++i) { + auto column = original->ColumnChunk(i); + auto chunk = group->NextColumnChunk(); + if (auto stats = column->statistics()) { + chunk->SetStatistics(stats->Encode()); + } + std::map<::parquet::Encoding::type, int32_t> dictionary_encodings; + std::map<::parquet::Encoding::type, int32_t> data_encodings; + if (encoding_stats) { + for (const auto& stat : column->encoding_stats()) { + if (stat.page_type == ::parquet::PageType::DICTIONARY_PAGE) { + dictionary_encodings[stat.encoding] += stat.count; + } else { + data_encodings[stat.encoding] += stat.count; + } + } + } + chunk->Finish(column->num_values(), + dictionary_offset.value_or(column->dictionary_page_offset()), 0, + data_offset.value_or(column->data_page_offset()), + column->total_compressed_size(), column->total_uncompressed_size(), + column->has_dictionary_page(), false, dictionary_encodings, + data_encodings); + if (corrupt_stats && i == 0) { + // Reinitialize only the column metadata's physical type/path. The file + // schema remains unchanged, making statistics() throw on the mismatch. + (void)::parquet::ColumnChunkMetaDataBuilder::Make( + writer_properties_, metadata_->schema()->Column(1), + const_cast(chunk->contents())); + } + } + group->Finish(original->total_byte_size()); + } + return builder->Finish(metadata_->key_value_metadata()); + } + + void ReplaceFooter(const ::parquet::FileMetaData& metadata) { + auto mock = std::static_pointer_cast(io_); + auto& data = mock->FileData(path_); + data.resize(data.size() - metadata_->size() - 8); + auto out = ::arrow::io::BufferOutputStream::Create().ValueOrDie(); + ::parquet::WriteFileMetaData(metadata, out.get()); + auto footer = out->Finish().ValueOrDie(); + auto bytes = std::as_bytes(std::span(footer->data(), footer->size())); + data.insert(data.end(), bytes.begin(), bytes.end()); + } +}; + +TEST_F(ParquetDictionaryFilterTest, ComparisonsAndBooleanExpressions) { + ASSERT_THAT(Write(true, "[[0,0],[10,1],[5,2],[15,3],[0,4],[10,5]]"), IsOk()); + CheckDictionaries(Expressions::Equal("key", Literal::Int(5)), {false, true, false}); + CheckDictionaries(Expressions::NotEqual("key", Literal::Int(5)), {true, true, true}); + CheckDictionaries(Expressions::LessThan("key", Literal::Int(5)), {true, false, true}); + CheckDictionaries(Expressions::LessThanOrEqual("key", Literal::Int(5)), + {true, true, true}); + CheckDictionaries(Expressions::GreaterThan("key", Literal::Int(10)), + {false, true, false}); + CheckDictionaries(Expressions::GreaterThanOrEqual("key", Literal::Int(10)), + {true, true, true}); + auto equal = Expressions::Equal("key", Literal::Int(5)); + CheckDictionaries(Expressions::And(equal, Expressions::AlwaysTrue()), + {false, true, false}); + CheckDictionaries(Expressions::Or(equal, Expressions::AlwaysFalse()), + {false, true, false}); + CheckDictionaries(Expressions::Or(equal, Expressions::AlwaysTrue()), + {true, true, true}); + // Negating a bound range adds an IS_NULL branch (a required leaf can have + // an optional ancestor), which dictionary filtering leaves to statistics. + CheckDictionaries( + Expressions::Not(Expressions::GreaterThanOrEqual("key", Literal::Int(5))), + {true, true, true}); + CheckDictionaries(Expressions::Not(Expressions::NotEqual("key", Literal::Int(5))), + {false, true, false}); + Check( + Options(Expressions::Not(Expressions::GreaterThanOrEqual("key", Literal::Int(5)))), + {0, 1, 4, 5}); + // All three footer ranges contain 5. Only the second dictionary contains it. + Check(Options(equal), {2, 3}); + auto options = Options(equal); + options.properties.Set(ReaderProperties::kParquetRowGroupFilter, false); + Check(options, {0, 1, 2, 3, 4, 5}); +} + +TEST_F(ParquetDictionaryFilterTest, SetsNegativePredicatesAndNulls) { + ASSERT_THAT(Write(true, "[[5,0],[5,1],[null,2],[5,3],[0,4],[10,5]]"), IsOk()); + CheckDictionaries(Expressions::NotEqual("key", Literal::Int(5)), {false, true, true}); + CheckDictionaries(Expressions::In("key", {Literal::Int(5), Literal::Int(7)}), + {true, true, false}); + CheckDictionaries(Expressions::NotIn("key", {Literal::Int(5), Literal::Int(7)}), + {false, true, true}); + CheckDictionaries(Expressions::IsNull("key"), {true, true, true}); + CheckDictionaries(Expressions::NotNull("key"), {true, true, true}); + ASSERT_THAT(Write(false, "[[5,0],[5,1],[5,2],[5,3]]"), IsOk()); + CheckDictionaries(Expressions::NotEqual("key", Literal::Int(5)), {true, true}); + CheckDictionaries(Expressions::NotIn("key", {Literal::Int(5), Literal::Int(7)}), + {true, true}); +} + +TEST_F(ParquetDictionaryFilterTest, Prefixes) { + SetKeyType(string()); + ASSERT_THAT(Write(true, R"([["aa",0],["az",1],["ba",2],["bz",3],[null,4],["aa",5]])"), + IsOk()); + CheckDictionaries(Expressions::StartsWith("key", "a"), {true, false, true}); + CheckDictionaries(Expressions::NotStartsWith("key", "a"), {false, true, true}); + CheckDictionaries(Expressions::StartsWith("key", ""), {true, true, true}); + CheckDictionaries(Expressions::NotStartsWith("key", ""), {false, false, true}); +} + +TEST_F(ParquetDictionaryFilterTest, FloatingPointNaNAndSignedZero) { + for (const auto& type : std::vector>{float32(), float64()}) { + SetKeyType(type); + ASSERT_THAT( + Write(true, "[[NaN,0],[NaN,1],[null,2],[NaN,3],[0,4],[-0.0,5],[1,6],[NaN,7]]"), + IsOk()); + CheckDictionaries(Expressions::IsNaN("key"), {true, true, false, true}); + CheckDictionaries(Expressions::NotNaN("key"), {false, true, true, true}); + auto zero = + type->type_id() == TypeId::kFloat ? Literal::Float(0) : Literal::Double(0); + CheckDictionaries(Expressions::Equal("key", zero), {false, false, true, false}); + auto one = type->type_id() == TypeId::kFloat ? Literal::Float(1) : Literal::Double(1); + auto nan = type->type_id() == TypeId::kFloat + ? Literal::Float(std::numeric_limits::quiet_NaN()) + : Literal::Double(std::numeric_limits::quiet_NaN()); + CheckDictionaries(Expressions::In("key", {zero, one}), {false, false, true, true}); + CheckDictionaries(Expressions::NotIn("key", {zero, one}), {true, true, true, true}); + CheckDictionaries(Expressions::In("key", {nan, one}), {true, true, false, true}); + CheckDictionaries(Expressions::NotIn("key", {nan, one}), {false, true, true, false}); + Check(Options(Expressions::In("key", {nan, one})), {0, 1, 2, 3, 6, 7}); + Check(Options(Expressions::NotIn("key", {nan, one})), {2, 3, 4, 5}); + CheckDictionaries(Expressions::Not(Expressions::IsNaN("key")), + {false, true, true, true}); + } +} + +TEST_F(ParquetDictionaryFilterTest, PhysicalTypesAndPromotions) { + struct Case { + std::shared_ptr file_type; + std::shared_ptr read_type; + std::string json; + Literal value; + }; + const std::string numbers = "[[0,0],[10,1],[5,2],[15,3],[0,4],[10,5]]"; + const std::string decimals = + R"([["0.00",0],["10.00",1],["5.00",2],["15.00",3],["0.00",4],["10.00",5]])"; + for (const auto& test : + std::vector{{.file_type = int32(), + .read_type = int64(), + .json = numbers, + .value = Literal::Long(5)}, + {.file_type = int64(), + .read_type = int64(), + .json = numbers, + .value = Literal::Long(5)}, + {.file_type = float32(), + .read_type = float64(), + .json = numbers, + .value = Literal::Double(5)}, + {.file_type = float64(), + .read_type = float64(), + .json = numbers, + .value = Literal::Double(5)}, + {.file_type = date(), + .read_type = date(), + .json = numbers, + .value = Literal::Date(5)}, + {.file_type = binary(), + .read_type = binary(), + .json = R"([["a",0],["z",1],["m",2],["z",3],["a",4],["z",5]])", + .value = Literal::Binary({'m'})}, + {.file_type = fixed(1), + .read_type = fixed(1), + .json = R"([["a",0],["z",1],["m",2],["z",3],["a",4],["z",5]])", + .value = Literal::Fixed({'m'})}, + {.file_type = decimal(9, 2), + .read_type = decimal(18, 2), + .json = decimals, + .value = Literal::Decimal(500, 18, 2)}, + {.file_type = decimal(18, 2), + .read_type = decimal(18, 2), + .json = decimals, + .value = Literal::Decimal(500, 18, 2)}, + {.file_type = decimal(30, 2), + .read_type = decimal(30, 2), + .json = decimals, + .value = Literal::Decimal(500, 30, 2)}}) { + SCOPED_TRACE(test.file_type->ToString() + " -> " + test.read_type->ToString()); + SetKeyType(test.file_type); + ::parquet::WriterProperties::Builder properties; + writer_properties_ = + properties.enable_dictionary()->enable_store_decimal_as_integer()->build(); + ASSERT_THAT(Write(true, test.json), IsOk()); + SetKeyType(test.read_type); + CheckDictionaries(Expressions::Equal("key", test.value), {false, true, false}); + Check(Options(Expressions::Equal("key", test.value)), {2, 3}); + } +} + +TEST_F(ParquetDictionaryFilterTest, DataPageV2) { + ::parquet::WriterProperties::Builder properties; + writer_properties_ = properties.enable_dictionary() + ->data_page_version(::parquet::ParquetDataPageVersion::V2) + ->build(); + ASSERT_THAT(Write(false), IsOk()); + CheckDictionaries(Expressions::Equal("key", Literal::Int(10)), {false, true, false}); + Check(Options(Expressions::Equal("key", Literal::Int(10))), {2, 3}); +} + +TEST_F(ParquetDictionaryFilterTest, IncompatibleTypesAndMissingFields) { + ASSERT_THAT(Write(), IsOk()); + SetKeyType(date()); + CheckDictionaries(Expressions::Equal("key", Literal::Date(100)), {true, true, true}); + schema_ = std::make_shared( + std::vector{SchemaField::MakeOptional(3, "missing", int32())}); + CheckDictionaries(Expressions::Equal("missing", Literal::Int(100)), {true, true, true}); +} + +TEST_F(ParquetDictionaryFilterTest, NoDictionaryAndMixedEncodingRetainGroups) { + ::parquet::WriterProperties::Builder plain; + writer_properties_ = plain.disable_dictionary()->build(); + ASSERT_THAT(Write(), IsOk()); + CheckDictionaries(Expressions::Equal("key", Literal::Int(100)), {true, true, true}); + + ::parquet::WriterProperties::Builder fallback; + writer_properties_ = fallback.enable_dictionary() + ->dictionary_pagesize_limit(1) + ->write_batch_size(1) + ->data_pagesize(32) + ->build(); + row_group_size_ = 128; + ::arrow::Int32Builder keys; + ::arrow::Int64Builder positions; + for (int i = 0; i < row_group_size_; ++i) { + ASSERT_TRUE(keys.Append(2 * i).ok()); + ASSERT_TRUE(positions.Append(i).ok()); + } + ArrowSchema c_schema; + ASSERT_THAT(ToArrowSchema(*schema_, &c_schema), IsOk()); + auto arrow_schema = ::arrow::ImportSchema(&c_schema).ValueOrDie(); + auto table = ::arrow::Table::Make( + arrow_schema, {keys.Finish().ValueOrDie(), positions.Finish().ValueOrDie()}); + ASSERT_THAT(Write(table), IsOk()); + auto column = metadata_->RowGroup(0)->ColumnChunk(0); + ASSERT_TRUE(column->has_dictionary_page()); + ASSERT_TRUE(std::ranges::any_of(column->encoding_stats(), [](const auto& stat) { + return stat.page_type == ::parquet::PageType::DATA_PAGE && + stat.encoding == ::parquet::Encoding::PLAIN; + })); + CheckDictionaries(Expressions::Equal("key", Literal::Int(1)), {true}); + // Value 254 occurs after dictionary fallback and must remain readable. + auto options = Options(Expressions::Equal("key", Literal::Int(254))); + ICEBERG_UNWRAP_OR_FAIL(auto actual, Read(options)); + ASSERT_EQ(actual.size(), 128); + EXPECT_EQ(actual.back(), 127); +} + +TEST_F(ParquetDictionaryFilterTest, SplitPositionsRenameAndReuse) { + ASSERT_THAT(Write(true, "[[5,0],[15,1],[0,2],[10,3],[5,4],[15,5]]"), IsOk()); + auto filter = Expressions::Equal("key", Literal::Int(5)); + Check(Options(filter), {0, 1, 4, 5}); + auto options = Options(filter); + options.properties.Set(ReaderProperties::kBatchSize, int64_t{3}); + Check(options, {0, 1, 4, 5}); + ICEBERG_UNWRAP_OR_FAIL(auto input, arrow::OpenArrowInputStream(io_, path_)); + ICEBERG_UNWRAP_OR_FAIL(auto midpoint, RowGroupMidpoint(*metadata_->RowGroup(1), + input->GetSize().ValueOrDie())); + options.split = Split{.offset = static_cast(midpoint), .length = 1}; + Check(options, {}); + options.filter = Options(Expressions::Equal("key", Literal::Int(0))).filter; + Check(options, {2, 3}); + options.split->length = 0; + Check(options, {}); + + schema_ = std::make_shared( + std::vector{SchemaField::MakeOptional(1, "renamed", int32())}); + Check(Options(Expressions::Equal("RENAMED", Literal::Int(5)), false), {0, 1, 4, 5}); + ICEBERG_UNWRAP_OR_FAIL(auto reader, + ReaderFactoryRegistry::Open( + FileFormatType::kParquet, + Options(Expressions::Equal("renamed", Literal::Int(999))))); + ICEBERG_UNWRAP_OR_FAIL(auto empty, reader->Next()); + EXPECT_FALSE(empty.has_value()); + ASSERT_THAT(reader->Close(), IsOk()); + auto unfiltered = Options(nullptr); + ASSERT_THAT(reader->Open(unfiltered), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto batch, reader->Next()); + ASSERT_TRUE(batch.has_value()); + EXPECT_GT(batch->length, 0); + batch->release(&*batch); + ASSERT_THAT(reader->Close(), IsOk()); +} + +TEST_F(ParquetDictionaryFilterTest, NameMappingWithoutFieldIdsRetainsGroups) { + auto keys = ::arrow::json::ArrayFromJSONString(::arrow::int32(), "[0,1,10,11,20,21]") + .ValueOrDie(); + auto values = + ::arrow::json::ArrayFromJSONString(::arrow::int64(), "[0,1,2,3,4,5]").ValueOrDie(); + auto file_schema = ::arrow::schema({::arrow::field("stored_key", ::arrow::int32()), + ::arrow::field("stored_value", ::arrow::int64())}); + ASSERT_THAT(Write(::arrow::Table::Make(file_schema, {keys, values})), IsOk()); + ASSERT_EQ(metadata_->schema()->Column(0)->schema_node()->field_id(), -1); + auto mapping = + NameMapping::Make({MappedField{.names = {"stored_key"}, .field_id = 1}, + MappedField{.names = {"stored_value"}, .field_id = 2}}); + auto options = Options(Expressions::Equal("key", Literal::Int(100))); + options.name_mapping = std::move(mapping); + // Name mapping resolves the projection. Footer columns without IDs remain + // unavailable to statistics filtering, so the reader must keep their groups. + Check(options, {0, 1, 2, 3, 4, 5}); +} + +TEST_F(ParquetDictionaryFilterTest, UnsupportedTransformRetainsGroups) { + ASSERT_THAT(Write(), IsOk()); + Check(Options(Expressions::Equal(Expressions::Bucket("key", 16), + Literal::Int(15))), + {0, 1, 2, 3, 4, 5}); +} + +TEST_F(ParquetDictionaryFilterTest, MissingColumnWithDefaultRetainsGroups) { + ASSERT_THAT(Write(), IsOk()); + schema_ = std::make_shared(std::vector{ + schema_->fields()[0], schema_->fields()[1], + SchemaField::MakeOptional(3, "defaulted", int32()) + .WithInitialDefault(std::make_shared(Literal::Int(7)))}); + Check(Options(Expressions::Equal("defaulted", Literal::Int(7))), {0, 1, 2, 3, 4, 5}); +} + +TEST_F(ParquetDictionaryFilterTest, RequiredLeafWithNullableParent) { + SetKeyType(std::make_shared( + std::vector{SchemaField::MakeRequired(3, "nested", int32())})); + ASSERT_THAT(Write(true, "[[null,0],[null,1],[[10],2],[[11],3],[null,4],[null,5]]"), + IsOk()); + auto is_null = Expressions::IsNull("key.nested"); + auto not_null = Expressions::NotNull("key.nested"); + Check(Options(is_null), {0, 1, 4, 5}); + Check(Options(not_null), {2, 3}); + Check(Options(Expressions::Not(is_null)), {2, 3}); + Check(Options(Expressions::Not(not_null)), {0, 1, 4, 5}); + Check(Options(Expressions::Not(Expressions::LessThan("key.nested", Literal::Int(100)))), + {0, 1, 4, 5}); +} + +TEST_F(ParquetDictionaryFilterTest, NestedPrimitiveStatistics) { + SetKeyType(std::make_shared( + std::vector{SchemaField::MakeOptional(3, "nested", int32())})); + ASSERT_THAT(Write(true, "[[[0],0],[[1],1],[[10],2],[[11],3],[[20],4],[[21],5]]"), + IsOk()); + Check(Options(Expressions::Equal("key.nested", Literal::Int(100))), {}); + Check(Options(Expressions::Equal("key.nested", Literal::Int(10))), {2, 3}); +} + +TEST_F(ParquetDictionaryFilterTest, UuidBoundsUseUnsignedByteOrder) { + SetKeyType(uuid()); + // Group 0 holds UUIDs below 0x80..., group 1 holds UUIDs at or above it. A signed + // byte comparison would order group 1 before group 0. + std::vector uuids; + for (const auto* str : + {"00000000-0000-0000-0000-000000000001", "7fffffff-ffff-ffff-ffff-ffffffffffff", + "80000000-0000-0000-0000-000000000000", "ffffffff-ffff-ffff-ffff-fffffffffffe"}) { + ICEBERG_UNWRAP_OR_FAIL(auto uuid, Uuid::FromString(str)); + uuids.push_back(uuid); + } + ::arrow::FixedSizeBinaryBuilder keys(::arrow::fixed_size_binary(Uuid::kLength)); + for (const auto& uuid : uuids) { + ASSERT_TRUE(keys.Append(uuid.bytes().data()).ok()); + } + auto key_array = ::arrow::ExtensionType::WrapArray(::arrow::extension::uuid(), + keys.Finish().ValueOrDie()); + auto values = + ::arrow::json::ArrayFromJSONString(::arrow::int64(), "[0,1,2,3]").ValueOrDie(); + ArrowSchema c_schema; + ASSERT_THAT(ToArrowSchema(*schema_, &c_schema), IsOk()); + auto arrow_schema = ::arrow::ImportSchema(&c_schema).ValueOrDie(); + ASSERT_THAT(Write(::arrow::Table::Make(arrow_schema, {key_array, values})), IsOk()); + ASSERT_EQ(metadata_->num_row_groups(), 2); + ASSERT_TRUE(metadata_->RowGroup(0)->ColumnChunk(0)->statistics()->HasMinMax()); + + CheckDictionaries(Expressions::Equal("key", Literal::UUID(uuids[2])), {false, true}); + const auto& group0_max = uuids[1]; + const auto& group1_min = uuids[2]; + Check(Options(Expressions::LessThan("key", Literal::UUID(group1_min))), {0, 1}); + Check(Options(Expressions::LessThanOrEqual("key", Literal::UUID(group0_max))), {0, 1}); + Check(Options(Expressions::GreaterThan("key", Literal::UUID(group0_max))), {2, 3}); + Check(Options(Expressions::GreaterThanOrEqual("key", Literal::UUID(group1_min))), + {2, 3}); + Check(Options(Expressions::Equal("key", Literal::UUID(group1_min))), {2, 3}); + Check( + Options(Expressions::Not(Expressions::LessThan("key", Literal::UUID(group1_min)))), + {2, 3}); +} + +TEST_F(ParquetDictionaryFilterTest, EncodingMetadataAndDictionaryCache) { + ASSERT_THAT(Write(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto input, arrow::OpenArrowInputStream(io_, path_)); + ::parquet::arrow::SchemaManifest manifest; + ASSERT_TRUE(::parquet::arrow::SchemaManifest::Make( + metadata_->schema(), nullptr, + ::parquet::default_arrow_reader_properties(), &manifest) + .ok()); + auto expr = Expressions::And(Expressions::Equal("key", Literal::Int(100)), + Expressions::Equal("key", Literal::Int(101))); + ICEBERG_UNWRAP_OR_FAIL(auto filter, ParquetDictionaryRowGroupFilter::Make( + Options(expr).filter, *metadata_->schema())); + auto io = std::static_pointer_cast(io_); + io->reads.clear(); + ICEBERG_UNWRAP_OR_FAIL(auto result, + filter->ShouldRead(manifest, *metadata_, 0, input, + ::parquet::default_reader_properties())); + EXPECT_FALSE(result); + EXPECT_EQ(io->reads.size(), 1); // Both predicates share one dictionary read. + + auto rebuilt = RebuildMetadata(false, false); + ASSERT_TRUE(rebuilt->RowGroup(0)->ColumnChunk(0)->encoding_stats().empty()); + io->reads.clear(); + ICEBERG_UNWRAP_OR_FAIL(auto should_read, + filter->ShouldRead(manifest, *rebuilt, 0, input, + ::parquet::default_reader_properties())); + EXPECT_TRUE(should_read); + EXPECT_TRUE(io->reads.empty()); +} + +TEST_F(ParquetDictionaryFilterTest, MissingAndTruncatedDictionaryReturnErrors) { + ASSERT_THAT(Write(), IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto input, arrow::OpenArrowInputStream(io_, path_)); + ::parquet::arrow::SchemaManifest manifest; + ASSERT_TRUE(::parquet::arrow::SchemaManifest::Make( + metadata_->schema(), nullptr, + ::parquet::default_arrow_reader_properties(), &manifest) + .ok()); + ICEBERG_UNWRAP_OR_FAIL(auto filter, + ParquetDictionaryRowGroupFilter::Make( + Options(Expressions::Equal("key", Literal::Int(100))).filter, + *metadata_->schema())); + auto column = metadata_->RowGroup(0)->ColumnChunk(0); + int64_t start = column->dictionary_page_offset(); + int64_t end = column->data_page_offset(); + int64_t size = input->GetSize().ValueOrDie(); + for (const auto& range : + std::vector{{.offset = -1, .length = end}, + {.offset = 0, .length = end}, + {.offset = start, .length = start}, + {.offset = end, .length = start}, + {.offset = start, .length = size + 1}, + {.offset = start, .length = start + 1}, + {.offset = start, .length = end - 1}, + {.offset = end, .length = end + 1}}) { + SCOPED_TRACE(std::to_string(range.offset) + ":" + std::to_string(range.length)); + auto rebuilt = RebuildMetadata(false, true, range.offset, range.length); + EXPECT_FALSE(filter + ->ShouldRead(manifest, *rebuilt, 0, input, + ::parquet::default_reader_properties()) + .has_value()); + } + // Corrupt the page header without altering valid offsets. + auto& data = std::static_pointer_cast(io_)->FileData(path_); + std::ranges::fill(std::span(data).subspan(start, end - start), std::byte{0}); + EXPECT_FALSE(filter + ->ShouldRead(manifest, *metadata_, 0, input, + ::parquet::default_reader_properties()) + .has_value()); +} + +TEST_F(ParquetDictionaryFilterTest, OpenReadsOnlyDictionaryPage) { + row_group_size_ = 3 * 1024 * 1024; + ::parquet::WriterProperties::Builder properties; + writer_properties_ = properties.enable_dictionary() + ->compression(::parquet::Compression::UNCOMPRESSED) + ->data_pagesize(8 * 1024 * 1024) + ->max_row_group_length(row_group_size_) + ->max_rows_per_page(row_group_size_) + ->build(); + ::arrow::Int32Builder keys; + ::arrow::Int64Builder positions; + ASSERT_TRUE(keys.Reserve(row_group_size_).ok()); + ASSERT_TRUE(positions.Reserve(row_group_size_).ok()); + for (int64_t i = 0; i < row_group_size_; ++i) { + keys.UnsafeAppend(static_cast(2 * (i % 256))); + positions.UnsafeAppend(i); + } + ArrowSchema c_schema; + ASSERT_THAT(ToArrowSchema(*schema_, &c_schema), IsOk()); + auto arrow_schema = ::arrow::ImportSchema(&c_schema).ValueOrDie(); + ASSERT_THAT( + Write(::arrow::Table::Make( + arrow_schema, {keys.Finish().ValueOrDie(), positions.Finish().ValueOrDie()})), + IsOk()); + auto column = metadata_->RowGroup(0)->ColumnChunk(0); + ASSERT_EQ(metadata_->num_row_groups(), 1); + ASSERT_GT(column->total_compressed_size(), 2 * 1024 * 1024); + // Project only metadata columns so eager reader initialization cannot prefetch + // output data. This isolates the I/O needed for dictionary pruning. + projection_ = std::make_shared( + std::vector{MetadataColumns::kRowPosition, MetadataColumns::kRowId}); + auto io = std::static_pointer_cast(io_); + io->reads.clear(); + ICEBERG_UNWRAP_OR_FAIL( + auto unfiltered, + ReaderFactoryRegistry::Open(FileFormatType::kParquet, Options(nullptr))); + auto footer_reads = io->reads; + ASSERT_THAT(unfiltered->Close(), IsOk()); + for (int value : {1, 2}) { // Within footer bounds: absent/present in dictionary. + io->reads.clear(); + ICEBERG_UNWRAP_OR_FAIL(auto reader, + ReaderFactoryRegistry::Open( + FileFormatType::kParquet, + Options(Expressions::Equal("key", Literal::Int(value))))); + auto expected = footer_reads; + expected.push_back( + {.offset = column->dictionary_page_offset(), + .length = column->data_page_offset() - column->dictionary_page_offset()}); + EXPECT_EQ(io->reads, expected); + ASSERT_THAT(reader->Close(), IsOk()); + } +} + +TEST_F(ParquetDictionaryFilterTest, UnreadableStatisticsRetainNegativePredicates) { + struct Case { + std::shared_ptr type; + std::string json; + std::vector> predicates; + }; + for (const auto& test : + std::vector{{.type = int32(), + .json = "[[5,0],[5,1]]", + .predicates = {Expressions::NotEqual("key", Literal::Int(5)), + Expressions::NotIn( + "key", {Literal::Int(5), Literal::Int(6)})}}, + {.type = float64(), + .json = "[[NaN,0],[NaN,1]]", + .predicates = {Expressions::NotNaN("key")}}, + {.type = string(), + .json = R"([["abc",0],["abc",1]])", + .predicates = {Expressions::NotStartsWith("key", "a")}}}) { + SetKeyType(test.type); + ASSERT_THAT(Write(true, test.json), IsOk()); + auto rebuilt = RebuildMetadata(true); + EXPECT_THROW(rebuilt->RowGroup(0)->ColumnChunk(0)->statistics(), + ::parquet::ParquetException); + ::parquet::arrow::SchemaManifest manifest; + ASSERT_TRUE(::parquet::arrow::SchemaManifest::Make( + metadata_->schema(), nullptr, + ::parquet::default_arrow_reader_properties(), &manifest) + .ok()); + ICEBERG_UNWRAP_OR_FAIL(auto input, arrow::OpenArrowInputStream(io_, path_)); + for (const auto& predicate : test.predicates) { + SCOPED_TRACE(predicate->ToString()); + ICEBERG_UNWRAP_OR_FAIL( + auto filter, ParquetDictionaryRowGroupFilter::Make(Options(predicate).filter, + *metadata_->schema())); + ICEBERG_UNWRAP_OR_FAIL(auto clean_result, + filter->ShouldRead(manifest, *metadata_, 0, input, + ::parquet::default_reader_properties())); + EXPECT_FALSE(clean_result); + ICEBERG_UNWRAP_OR_FAIL(auto result, + filter->ShouldRead(manifest, *rebuilt, 0, input, + ::parquet::default_reader_properties())); + EXPECT_TRUE(result); + } + ReplaceFooter(*rebuilt); + for (const auto& predicate : test.predicates) { + Check(Options(predicate), {0, 1}); // Includes footer filtering and Open(). + } + } +} + +TEST_F(ParquetDictionaryFilterTest, CompressedDictionaryPages) { + if (!::arrow::util::Codec::IsAvailable(::arrow::Compression::GZIP)) { + GTEST_SKIP() << "Gzip is not available"; + } + for (auto version : + {::parquet::ParquetDataPageVersion::V1, ::parquet::ParquetDataPageVersion::V2}) { + ::parquet::WriterProperties::Builder properties; + writer_properties_ = properties.enable_dictionary() + ->compression(::parquet::Compression::GZIP) + ->data_page_version(version) + ->build(); + ASSERT_THAT(Write(false), IsOk()); + CheckDictionaries(Expressions::Equal("key", Literal::Int(10)), {false, true, false}); + Check(Options(Expressions::Equal("key", Literal::Int(10))), {2, 3}); + } +} + +TEST_F(ParquetDictionaryFilterTest, SignedZerosRemainDistinct) { + for (const auto& type : std::vector>{float32(), float64()}) { + SetKeyType(type); + ASSERT_THAT(Write(true, "[[-0.0,0],[-0.0,1]]"), IsOk()); + auto positive = + type->type_id() == TypeId::kFloat ? Literal::Float(0.0F) : Literal::Double(0.0); + auto negative = + type->type_id() == TypeId::kFloat ? Literal::Float(-0.0F) : Literal::Double(-0.0); + CheckDictionaries(Expressions::Equal("key", positive), {false}); + CheckDictionaries(Expressions::Equal("key", negative), {true}); + CheckDictionaries(Expressions::NotEqual("key", positive), {true}); + CheckDictionaries(Expressions::NotEqual("key", negative), {false}); + CheckDictionaries(Expressions::In("key", {positive, type->type_id() == TypeId::kFloat + ? Literal::Float(1) + : Literal::Double(1)}), + {false}); + Check(Options(Expressions::Equal("key", positive)), {}); + Check(Options(Expressions::Equal("key", negative)), {0, 1}); + } +} + } // namespace } // namespace iceberg::parquet