diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index 4cde4fb23..df6a0bfaf 100644 --- a/src/iceberg/CMakeLists.txt +++ b/src/iceberg/CMakeLists.txt @@ -58,6 +58,7 @@ set(ICEBERG_SOURCES geospatial.cc incremental_append_scan.cc incremental_changelog_scan.cc + inspect/partitions_table.cc inspect/history_table.cc inspect/metadata_scan.cc inspect/metadata_table.cc diff --git a/src/iceberg/cpp20_compatibility_internal.h b/src/iceberg/cpp20_compatibility_internal.h index d9c64d36a..d734b4df8 100644 --- a/src/iceberg/cpp20_compatibility_internal.h +++ b/src/iceberg/cpp20_compatibility_internal.h @@ -85,6 +85,7 @@ #include "iceberg/inspect/metadata_scan_task.h" #include "iceberg/inspect/metadata_scan_task_reader.h" #include "iceberg/inspect/metadata_table.h" +#include "iceberg/inspect/partitions_table.h" #include "iceberg/inspect/refs_table.h" #include "iceberg/inspect/snapshots_table.h" #include "iceberg/labels.h" diff --git a/src/iceberg/inspect/metadata_table.cc b/src/iceberg/inspect/metadata_table.cc index 012ffaa99..9c601cddc 100644 --- a/src/iceberg/inspect/metadata_table.cc +++ b/src/iceberg/inspect/metadata_table.cc @@ -24,6 +24,7 @@ #include "iceberg/inspect/history_table.h" #include "iceberg/inspect/metadata_scan.h" +#include "iceberg/inspect/partitions_table.h" #include "iceberg/inspect/refs_table.h" #include "iceberg/inspect/snapshots_table.h" @@ -43,6 +44,8 @@ Result> MetadataTable::Make( return RefsTable::Make(std::move(table)); case MetadataTableKind::kHistory: return HistoryTable::Make(std::move(table)); + case MetadataTableKind::kPartitions: + return PartitionsTable::Make(std::move(table)); default: break; } diff --git a/src/iceberg/inspect/partitions_table.cc b/src/iceberg/inspect/partitions_table.cc new file mode 100644 index 000000000..2501d5188 --- /dev/null +++ b/src/iceberg/inspect/partitions_table.cc @@ -0,0 +1,341 @@ +/* + * 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/inspect/partitions_table.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "iceberg/arrow_row_builder_internal.h" +#include "iceberg/expression/literal.h" +#include "iceberg/inspect/metadata_manifest_util_internal.h" +#include "iceberg/inspect/metadata_partition_util_internal.h" +#include "iceberg/inspect/metadata_scan_util_internal.h" +#include "iceberg/inspect/metadata_table_util_internal.h" +#include "iceberg/manifest/manifest_entry.h" +#include "iceberg/partition_spec.h" +#include "iceberg/row/partition_values.h" +#include "iceberg/schema.h" +#include "iceberg/schema_field.h" +#include "iceberg/snapshot.h" +#include "iceberg/table.h" +#include "iceberg/table_metadata.h" +#include "iceberg/type.h" +#include "iceberg/util/macros.h" + +namespace iceberg { +namespace { + +constexpr int32_t kPartitionFieldId = 1; +constexpr int32_t kRecordCountFieldId = 2; +constexpr int32_t kFileCountFieldId = 3; +constexpr int32_t kSpecIdFieldId = 4; +constexpr int32_t kPositionDeleteRecordCountFieldId = 5; +constexpr int32_t kPositionDeleteFileCountFieldId = 6; +constexpr int32_t kEqualityDeleteRecordCountFieldId = 7; +constexpr int32_t kEqualityDeleteFileCountFieldId = 8; +constexpr int32_t kLastUpdatedAtFieldId = 9; +constexpr int32_t kLastUpdatedSnapshotIdFieldId = 10; +constexpr int32_t kTotalDataFileSizeFieldId = 11; + +struct PartitionKey { + PartitionValues values; + size_t projected_fields; + + bool operator==(const PartitionKey& other) const { + if (projected_fields != other.projected_fields || + values.num_fields() != other.values.num_fields()) { + return false; + } + for (size_t index = 0; index < values.num_fields(); ++index) { + const auto& lhs = values.values()[index]; + const auto& rhs = other.values.values()[index]; + if (lhs.IsNull() || rhs.IsNull()) { + if (lhs.IsNull() != rhs.IsNull()) { + return false; + } + } else if (lhs.IsNaN() && rhs.IsNaN()) { + // Java partition equality canonicalizes all NaN signs and payloads. + continue; + } else if (lhs != rhs) { + return false; + } + } + return true; + } +}; + +struct PartitionKeyHash { + size_t operator()(const PartitionKey& key) const noexcept { + size_t result = 17; + for (const auto& value : key.values.values()) { + const size_t value_hash = + value.IsNaN() ? 0x7ff8000000000000ULL : LiteralHash{}(value); + result = result * 37 + value_hash; + } + return result * 37 + key.projected_fields; + } +}; + +size_t ProjectedFieldCount(const StructType& partition_type, const PartitionSpec& spec) { + size_t count = 0; + for (const auto& field : partition_type.fields()) { + count += std::ranges::any_of( + spec.fields(), [field_id = field.field_id()](const PartitionField& spec_field) { + return spec_field.field_id() == field_id; + }); + } + return count; +} + +struct PartitionStats { + explicit PartitionStats(PartitionValues values) : partition(std::move(values)) {} + + PartitionValues partition; + int32_t spec_id = PartitionSpec::kInitialSpecId; + int64_t data_record_count = 0; + int32_t data_file_count = 0; + int64_t data_file_size = 0; + int64_t position_delete_record_count = 0; + int32_t position_delete_file_count = 0; + int64_t equality_delete_record_count = 0; + int32_t equality_delete_file_count = 0; + std::optional last_updated_at; + std::optional last_updated_snapshot_id; +}; + +Status AppendPartition(ArrowRowBuilder& builder, const Schema& schema, + const StructType& partition_type, + const PartitionStats& partition) { + for (size_t index = 0; index < schema.fields().size(); ++index) { + auto* array = builder.column(index); + switch (schema.fields()[index].field_id()) { + case kPartitionFieldId: + ICEBERG_RETURN_UNEXPECTED( + internal::AppendPartitionValues(array, partition_type, partition.partition)); + break; + case kSpecIdFieldId: + ICEBERG_RETURN_UNEXPECTED(AppendInt(array, partition.spec_id)); + break; + case kRecordCountFieldId: + ICEBERG_RETURN_UNEXPECTED(AppendInt(array, partition.data_record_count)); + break; + case kFileCountFieldId: + ICEBERG_RETURN_UNEXPECTED(AppendInt(array, partition.data_file_count)); + break; + case kTotalDataFileSizeFieldId: + ICEBERG_RETURN_UNEXPECTED(AppendInt(array, partition.data_file_size)); + break; + case kPositionDeleteRecordCountFieldId: + ICEBERG_RETURN_UNEXPECTED( + AppendInt(array, partition.position_delete_record_count)); + break; + case kPositionDeleteFileCountFieldId: + ICEBERG_RETURN_UNEXPECTED(AppendInt(array, partition.position_delete_file_count)); + break; + case kEqualityDeleteRecordCountFieldId: + ICEBERG_RETURN_UNEXPECTED( + AppendInt(array, partition.equality_delete_record_count)); + break; + case kEqualityDeleteFileCountFieldId: + ICEBERG_RETURN_UNEXPECTED(AppendInt(array, partition.equality_delete_file_count)); + break; + case kLastUpdatedAtFieldId: + if (partition.last_updated_at.has_value()) { + ICEBERG_RETURN_UNEXPECTED( + AppendInt(array, std::chrono::duration_cast( + partition.last_updated_at->time_since_epoch()) + .count())); + } else { + ICEBERG_RETURN_UNEXPECTED(AppendNull(array)); + } + break; + case kLastUpdatedSnapshotIdFieldId: + if (partition.last_updated_snapshot_id.has_value()) { + ICEBERG_RETURN_UNEXPECTED( + AppendInt(array, *partition.last_updated_snapshot_id)); + } else { + ICEBERG_RETURN_UNEXPECTED(AppendNull(array)); + } + break; + default: + return InvalidSchema("Unsupported partitions metadata field {}", + schema.fields()[index].field_id()); + } + } + return builder.FinishRow(); +} + +void UpdateCounts(PartitionStats& partition, const DataFile& file) { + switch (file.content) { + case DataFile::Content::kData: + partition.data_record_count += file.record_count; + ++partition.data_file_count; + partition.data_file_size += file.file_size_in_bytes; + break; + case DataFile::Content::kPositionDeletes: + partition.position_delete_record_count += file.record_count; + ++partition.position_delete_file_count; + break; + case DataFile::Content::kEqualityDeletes: + partition.equality_delete_record_count += file.record_count; + ++partition.equality_delete_file_count; + break; + } +} + +} // namespace + +PartitionsTable::PartitionsTable(std::shared_ptr table) + : MetadataTable(std::move(table)) {} + +PartitionsTable::~PartitionsTable() = default; + +std::shared_ptr PartitionsTable::schema() const { + std::lock_guard lock(schema_mutex_); + return schema_; +} + +Result> PartitionsTable::Make( + std::shared_ptr table) { + ICEBERG_PRECHECK(table != nullptr, "Table cannot be null"); + auto partitions = + std::unique_ptr(new PartitionsTable(std::move(table))); + ICEBERG_RETURN_UNEXPECTED( + partitions->SchemaForScan(*partitions->source_table()->metadata())); + return partitions; +} + +Result> PartitionsTable::SchemaForScan( + const TableMetadata& metadata) const { + ICEBERG_ASSIGN_OR_RAISE(auto partition_type, internal::UnifiedPartitionType(metadata)); + + std::vector fields; + if (!partition_type->fields().empty()) { + fields.push_back( + SchemaField::MakeRequired(kPartitionFieldId, "partition", partition_type)); + fields.push_back(SchemaField::MakeRequired(kSpecIdFieldId, "spec_id", int32())); + } + fields.push_back( + SchemaField::MakeRequired(kRecordCountFieldId, "record_count", int64())); + fields.push_back(SchemaField::MakeRequired(kFileCountFieldId, "file_count", int32())); + fields.push_back(SchemaField::MakeRequired(kTotalDataFileSizeFieldId, + "total_data_file_size_in_bytes", int64())); + fields.push_back(SchemaField::MakeRequired(kPositionDeleteRecordCountFieldId, + "position_delete_record_count", int64())); + fields.push_back(SchemaField::MakeRequired(kPositionDeleteFileCountFieldId, + "position_delete_file_count", int32())); + fields.push_back(SchemaField::MakeRequired(kEqualityDeleteRecordCountFieldId, + "equality_delete_record_count", int64())); + fields.push_back(SchemaField::MakeRequired(kEqualityDeleteFileCountFieldId, + "equality_delete_file_count", int32())); + fields.push_back(SchemaField::MakeOptional(kLastUpdatedAtFieldId, "last_updated_at", + timestamp_tz())); + fields.push_back(SchemaField::MakeOptional(kLastUpdatedSnapshotIdFieldId, + "last_updated_snapshot_id", int64())); + + auto schema = std::make_shared(std::move(fields)); + { + std::lock_guard lock(schema_mutex_); + schema_ = schema; + } + return schema; +} + +Result> PartitionsTable::MakeScan( + internal::MetadataScanContext context) const { + return internal::MakeSnapshotMetadataScan(std::move(context), ReadRows); +} + +Result PartitionsTable::ReadRows( + const internal::MetadataScanContext& context, const std::shared_ptr& io, + int64_t batch_size) { + ICEBERG_ASSIGN_OR_RAISE(auto partition_type, + internal::UnifiedPartitionType(*context.table_metadata)); + ICEBERG_ASSIGN_OR_RAISE(auto snapshot, internal::ResolveMetadataTableSnapshot( + *context.table_metadata, context)); + ICEBERG_ASSIGN_OR_RAISE( + auto files, internal::LiveFiles( + *context.table_metadata, io, snapshot, + {"content", "partition", "record_count", "file_size_in_bytes"})); + + std::vector partitions; + std::unordered_map positions; + std::unordered_map> snapshots; + while (true) { + ICEBERG_ASSIGN_OR_RAISE(auto next_file, files->Next()); + if (!next_file.has_value()) { + break; + } + const auto& live_file = *next_file; + ICEBERG_ASSIGN_OR_RAISE(auto partition_values, internal::ProjectPartitionValues( + *partition_type, *live_file.spec, + live_file.file->partition)); + PartitionKey key{ + .values = std::move(partition_values), + .projected_fields = ProjectedFieldCount(*partition_type, *live_file.spec)}; + auto [position, inserted] = positions.try_emplace(key, partitions.size()); + if (inserted) { + partitions.emplace_back(std::move(key.values)); + } + auto& partition = partitions[position->second]; + UpdateCounts(partition, *live_file.file); + + if (live_file.snapshot_id.has_value()) { + auto snapshot_iter = snapshots.find(*live_file.snapshot_id); + if (snapshot_iter == snapshots.end()) { + auto file_snapshot = context.table_metadata->SnapshotById(*live_file.snapshot_id); + if (!file_snapshot.has_value() && + file_snapshot.error().kind != ErrorKind::kNotFound) { + return unexpected(file_snapshot.error()); + } + snapshot_iter = + snapshots.emplace(*live_file.snapshot_id, file_snapshot.value_or(nullptr)) + .first; + } + const auto& file_snapshot = snapshot_iter->second; + if (file_snapshot != nullptr && + (!partition.last_updated_at.has_value() || + file_snapshot->timestamp_ms > *partition.last_updated_at)) { + partition.spec_id = live_file.spec->spec_id(); + partition.last_updated_at = file_snapshot->timestamp_ms; + partition.last_updated_snapshot_id = file_snapshot->snapshot_id; + } + } + } + + auto schema = context.projected_schema; + return internal::MakeMetadataBatchStream( + schema, std::move(partitions), batch_size, + [schema, partition_type = std::move(partition_type)]( + ArrowRowBuilder& builder, const PartitionStats& partition) { + return AppendPartition(builder, *schema, *partition_type, partition); + }); +} + +} // namespace iceberg diff --git a/src/iceberg/inspect/partitions_table.h b/src/iceberg/inspect/partitions_table.h new file mode 100644 index 000000000..89f3652ea --- /dev/null +++ b/src/iceberg/inspect/partitions_table.h @@ -0,0 +1,68 @@ +/* + * 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 + +/// \file iceberg/inspect/partitions_table.h +/// \brief Define the partitions metadata table. + +#include +#include + +#include "iceberg/arrow_c_data.h" +#include "iceberg/iceberg_export.h" +#include "iceberg/inspect/metadata_table.h" +#include "iceberg/result.h" +#include "iceberg/type_fwd.h" + +namespace iceberg { + +/// \brief Metadata table containing aggregate file statistics by partition. +class ICEBERG_EXPORT PartitionsTable : public MetadataTable { + public: + static Result> Make( + std::shared_ptr table); + + ~PartitionsTable() override; + + Kind kind() const noexcept override { return Kind::kPartitions; } + + /// \brief Schema from construction or the most recent scan builder. + /// Existing scans and tasks retain their captured schema. + std::shared_ptr schema() const override; + + protected: + Result> SchemaForScan( + const TableMetadata& metadata) const override; + + Result> MakeScan( + internal::MetadataScanContext context) const override; + + private: + static Result ReadRows(const internal::MetadataScanContext& context, + const std::shared_ptr& io, + int64_t batch_size); + + explicit PartitionsTable(std::shared_ptr table); + + mutable std::mutex schema_mutex_; + mutable std::shared_ptr schema_; +}; + +} // namespace iceberg diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 1874fc25e..bd9aa3f70 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -209,7 +209,8 @@ if(ICEBERG_BUILD_BUNDLE) inspect/history_table_test.cc inspect/metadata_table_test.cc inspect/refs_table_test.cc - inspect/snapshots_table_test.cc) + inspect/snapshots_table_test.cc + inspect/partitions_table_test.cc) add_iceberg_test(eval_expr_test USE_BUNDLE diff --git a/src/iceberg/test/inspect/metadata_table_scan_test_base.h b/src/iceberg/test/inspect/metadata_table_scan_test_base.h new file mode 100644 index 000000000..218fdc325 --- /dev/null +++ b/src/iceberg/test/inspect/metadata_table_scan_test_base.h @@ -0,0 +1,270 @@ +/* + * 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 +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "iceberg/inspect/metadata_scan.h" +#include "iceberg/inspect/metadata_scan_task_reader.h" +#include "iceberg/inspect/metadata_table.h" +#include "iceberg/row/partition_values.h" +#include "iceberg/table_identifier.h" +#include "iceberg/test/mock_catalog.h" +#include "iceberg/test/scan_test_base.h" +#include "iceberg/util/conversions.h" + +namespace iceberg { + +struct SnapshotSelection { + std::variant snapshot; + std::string ref_name; +}; + +class MetadataTableScanTestBase : public ScanTestBase { + protected: + void CheckConcurrentBuilders(MetadataTableKind kind) { + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({}, kInvalidSnapshotId)); + ICEBERG_UNWRAP_OR_FAIL(auto metadata_table, table->Inspect(kind)); + const auto held_schema = metadata_table->schema(); + const size_t num_fields = held_schema->fields().size(); + std::atomic failures = 0; + std::barrier start(4); + std::vector threads; + for (int thread = 0; thread < 4; ++thread) { + threads.emplace_back([&] { + start.arrive_and_wait(); + for (int iteration = 0; iteration < 1000; ++iteration) { + auto builder = iteration % 2 == 0 ? metadata_table->NewScan() + : MetadataScanBuilder::Make(*metadata_table); + if (!builder || !(*builder)->Build()) { + ++failures; + } + auto schema = metadata_table->schema(); + if (!schema || schema->fields().size() != num_fields || + held_schema->fields().size() != num_fields) { + ++failures; + } + } + }); + } + for (auto& thread : threads) { + thread.join(); + } + EXPECT_EQ(failures.load(), 0); + } + + void CheckTaskProjectionAndLifetime(MetadataTableKind kind) { + std::vector manifests; + for (int32_t index = 0; index < 3; ++index) { + manifests.push_back(WriteDataManifest( + GetParam(), 10, + {MakeEntry( + ManifestStatus::kAdded, 10, 1, + MakeDataFile(std::format("file-{}.parquet", index), + PartitionValues(Literal::Int(index)), partitioned_spec_))}, + partitioned_spec_)); + } + auto first = std::make_shared( + Snapshot{.snapshot_id = 10, + .sequence_number = 1, + .timestamp_ms = TimePointMsFromUnixMs(1000), + .manifest_list = WriteManifestList(GetParam(), 10, 0, 1, manifests)}); + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({first}, 10, {}, partitioned_spec_)); + ICEBERG_UNWRAP_OR_FAIL(auto metadata_table, table->Inspect(kind)); + const auto canonical_schema = metadata_table->schema(); + const auto& fields = canonical_schema->fields(); + auto projection = + std::make_shared(std::vector{fields[1], fields[0]}); + const std::string first_column(projection->fields()[0].name()); + const std::string second_column(projection->fields()[1].name()); + ICEBERG_UNWRAP_OR_FAIL(auto builder, metadata_table->NewScan()); + ICEBERG_UNWRAP_OR_FAIL(auto empty_builder, metadata_table->NewScan()); + ICEBERG_UNWRAP_OR_FAIL(auto invalid_builder, metadata_table->NewScan()); + auto invalid_projection = std::make_shared(std::vector{ + SchemaField::MakeRequired(fields[0].field_id(), "wrong_type", int64())}); + EXPECT_FALSE(invalid_builder->Project(invalid_projection).Build().has_value()); + + // Builders retain metadata even when the source refreshes before Build(). + ASSERT_THAT(RefreshTableSchemaAndPartitionSpec(table, first, arrow::MakeMockFileIO()), + IsOk()); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Project(projection).Build()); + ICEBERG_UNWRAP_OR_FAIL(auto empty_scan, empty_builder->Select({}).Build()); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanTasks()); + ICEBERG_UNWRAP_OR_FAIL(auto empty_tasks, empty_scan->PlanTasks()); + ASSERT_EQ(tasks.size(), 1); + ASSERT_EQ(empty_tasks.size(), 1); + builder.reset(); + empty_builder.reset(); + invalid_builder.reset(); + scan.reset(); + empty_scan.reset(); + metadata_table.reset(); + table.reset(); + + ICEBERG_UNWRAP_OR_FAIL(auto reader, MetadataScanTaskReader::Make({.batch_size = 2})); + ICEBERG_UNWRAP_OR_FAIL(auto stream, reader->Open(*tasks[0])); + ICEBERG_UNWRAP_OR_FAIL(auto repeat_stream, reader->Open(*tasks[0])); + ICEBERG_UNWRAP_OR_FAIL(auto empty_stream, reader->Open(*empty_tasks[0])); + reader.reset(); + tasks.clear(); + empty_tasks.clear(); + for (auto* output : {&stream, &repeat_stream}) { + ICEBERG_UNWRAP_OR_FAIL(auto batches, ReadAllBatches(std::move(*output))); + ASSERT_EQ(batches.size(), 2); + EXPECT_EQ(batches[0]->num_rows(), 2); + EXPECT_EQ(batches[1]->num_rows(), 1); + ASSERT_EQ(batches[0]->num_columns(), 2); + EXPECT_EQ(batches[0]->schema()->field(0)->name(), first_column); + EXPECT_EQ(batches[0]->schema()->field(1)->name(), second_column); + } + ICEBERG_UNWRAP_OR_FAIL(auto empty_batches, ReadAllBatches(std::move(empty_stream))); + ASSERT_EQ(empty_batches.size(), 2); + EXPECT_EQ(empty_batches[0]->num_rows(), 2); + EXPECT_EQ(empty_batches[1]->num_rows(), 1); + EXPECT_EQ(empty_batches[0]->num_columns(), 0); + } + + template + static Result Scan(const std::unique_ptr& table, + const SnapshotSelection& selection = {}) { + ICEBERG_ASSIGN_OR_RAISE(auto builder, table->NewScan()); + if (auto id = std::get_if(&selection.snapshot)) { + builder->UseSnapshot(*id); + } else if (auto time = std::get_if(&selection.snapshot)) { + builder->AsOfTime(time->time_since_epoch().count()); + } + if (!selection.ref_name.empty()) { + builder->UseRef(selection.ref_name); + } + ICEBERG_ASSIGN_OR_RAISE(auto scan, builder->Build()); + return scan->ToArrowArrayStream(); + } + + void SetUp() override { + ScanTestBase::SetUp(); + catalog_ = std::make_shared(); + } + + Result> MakeTable( + std::vector> snapshots, int64_t current_snapshot_id, + std::unordered_map> refs = {}, + std::shared_ptr spec = nullptr) { + auto metadata = MakeTableMetadata(snapshots, current_snapshot_id, refs, spec); + return Table::Make( + TableIdentifier{.ns = Namespace{.levels = {"db"}}, .name = "table"}, + std::move(metadata), "s3://bucket/metadata.json", file_io_, catalog_); + } + + static Result>> ReadAllBatches( + ArrowArrayStream&& stream) { + auto reader = ::arrow::ImportRecordBatchReader(&stream); + if (!reader.ok()) { + return InvalidArrowData(reader.status().ToString()); + } + auto batches = reader.ValueUnsafe()->ToRecordBatches(); + if (!batches.ok()) { + return InvalidArrowData(batches.status().ToString()); + } + return std::move(batches).MoveValueUnsafe(); + } + + Status RefreshTableSchemaAndPartitionSpec( + const std::shared_ptr& table, const std::shared_ptr& first, + std::shared_ptr refreshed_io = nullptr) { + auto old_schema = schema_; + schema_ = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeRequired(2, "data", string()), + SchemaField::MakeOptional(3, "extra", int64())}, + 1); + std::shared_ptr new_spec; + ICEBERG_ASSIGN_OR_RAISE( + new_spec, + PartitionSpec::Make( + 2, {PartitionField(2, 1000, "data_bucket_16_2", Transform::Bucket(16)), + PartitionField(1, 1001, "id", Transform::Identity())})); + auto second = MakeAppendSnapshotWithPartitionValues( + GetParam(), 20, 10, 2, + {{"one.parquet", PartitionValues({Literal::Int(7), Literal::Int(1)})}, + {"two.parquet", PartitionValues({Literal::Int(7), Literal::Int(2)})}}, + new_spec); + auto metadata = MakeTableMetadata({first, second}, 20, {}, new_spec); + metadata->schemas.insert(metadata->schemas.begin(), old_schema); + metadata->partition_specs.push_back(new_spec); + metadata->last_column_id = 3; + metadata->last_partition_id = 1001; + ICEBERG_ASSIGN_OR_RAISE( + auto refreshed, + Table::Make(table->name(), std::move(metadata), "s3://bucket/refreshed.json", + refreshed_io != nullptr ? std::move(refreshed_io) : file_io_, + catalog_)); + EXPECT_CALL(*catalog_, LoadTable(::testing::_)) + .WillOnce(::testing::Return(refreshed)); + return table->Refresh(); + } + + Result> MakeDecimalPartitionTable( + const std::vector& values) { + schema_ = std::make_shared( + std::vector{SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeRequired(2, "amount", decimal(38, 2))}); + ICEBERG_ASSIGN_OR_RAISE( + partitioned_spec_, + PartitionSpec::Make(1, + {PartitionField(2, 1000, "amount", Transform::Identity())})); + std::vector entries; + for (size_t index = 0; index < values.size(); ++index) { + auto literal = Literal::Decimal(values[index], 38, 2); + auto file = MakeDataFile(std::format("decimal-{}.parquet", index), + PartitionValues(literal), partitioned_spec_); + ICEBERG_ASSIGN_OR_RAISE(auto bound, Conversions::ToBytes(literal)); + file->lower_bounds.emplace(2, bound); + file->upper_bounds.emplace(2, bound); + entries.push_back(MakeEntry(ManifestStatus::kAdded, 10, 1, std::move(file))); + } + auto manifest = + WriteDataManifest(GetParam(), 10, std::move(entries), partitioned_spec_); + auto snapshot = std::make_shared(Snapshot{ + .snapshot_id = 10, + .sequence_number = 1, + .timestamp_ms = TimePointMsFromUnixMs(1000), + .manifest_list = WriteManifestList(GetParam(), 10, 0, 1, {manifest}), + }); + return MakeTable({snapshot}, 10, {}, partitioned_spec_); + } + + std::shared_ptr catalog_; +}; + +} // namespace iceberg diff --git a/src/iceberg/test/inspect/partitions_table_test.cc b/src/iceberg/test/inspect/partitions_table_test.cc new file mode 100644 index 000000000..a0b15bf66 --- /dev/null +++ b/src/iceberg/test/inspect/partitions_table_test.cc @@ -0,0 +1,311 @@ +/* + * 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/inspect/partitions_table.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "iceberg/constants.h" +#include "iceberg/inspect/metadata_manifest_util_internal.h" +#include "iceberg/inspect/metadata_table.h" +#include "iceberg/partition_spec.h" +#include "iceberg/row/partition_values.h" +#include "iceberg/snapshot.h" +#include "iceberg/table.h" +#include "iceberg/table_identifier.h" +#include "iceberg/test/inspect/metadata_table_scan_test_base.h" +#include "iceberg/test/matchers.h" +#include "iceberg/transform.h" +#include "iceberg/util/conversions.h" +#include "iceberg/util/macros.h" + +namespace iceberg { +namespace { + +class PartitionsTableTest : public MetadataTableScanTestBase {}; + +TEST_P(PartitionsTableTest, Scan) { + auto snapshot = MakeAppendSnapshotWithPartitionValues( + GetParam(), 10, std::nullopt, 1, + {{"s3://bucket/data.parquet", PartitionValues(Literal::Int(7))}}, + partitioned_spec_); + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({snapshot}, 10, {}, partitioned_spec_)); + + ICEBERG_UNWRAP_OR_FAIL(auto partitions, MetadataTable::Make(table)); + ICEBERG_UNWRAP_OR_FAIL(auto partitions_stream, Scan(partitions)); + ICEBERG_UNWRAP_OR_FAIL(auto partition_batches, + ReadAllBatches(std::move(partitions_stream))); + ASSERT_EQ(partition_batches.size(), 1); + ASSERT_EQ(partition_batches[0]->num_rows(), 1); + auto records = std::static_pointer_cast<::arrow::Int64Array>( + partition_batches[0]->GetColumnByName("record_count")); + auto file_counts = std::static_pointer_cast<::arrow::Int32Array>( + partition_batches[0]->GetColumnByName("file_count")); + auto updated_snapshot_ids = std::static_pointer_cast<::arrow::Int64Array>( + partition_batches[0]->GetColumnByName("last_updated_snapshot_id")); + EXPECT_EQ(records->Value(0), 1); + EXPECT_EQ(file_counts->Value(0), 1); + EXPECT_EQ(updated_snapshot_ids->Value(0), 10); +} + +TEST_P(PartitionsTableTest, RejectsCombiningNamedRefsWithTimeTravel) { + auto first = MakeAppendSnapshot(GetParam(), 1, std::nullopt, 1, {"first.parquet"}); + auto second = MakeAppendSnapshot(GetParam(), 2, 1, 2, {"second.parquet"}); + ICEBERG_UNWRAP_OR_FAIL(auto branch, SnapshotRef::MakeBranch(1)); + ICEBERG_UNWRAP_OR_FAIL(auto tag, SnapshotRef::MakeTag(1)); + std::unordered_map> refs; + refs.emplace("dev", std::move(branch)); + refs.emplace("release", std::move(tag)); + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({first, second}, 2, std::move(refs))); + ICEBERG_UNWRAP_OR_FAIL(auto metadata_table, + MetadataTable::Make(table)); + for (const std::string ref_name : {"dev", "release"}) { + for (const SnapshotSelection selection : + {SnapshotSelection{.snapshot = int64_t{1}, .ref_name = ref_name}, + SnapshotSelection{.snapshot = first->timestamp_ms, .ref_name = ref_name}}) { + auto result = Scan(metadata_table, selection); + ASSERT_FALSE(result.has_value()); + EXPECT_EQ(result.error().kind, ErrorKind::kInvalidArgument); + } + } +} + +TEST_P(PartitionsTableTest, RejectsUnknownReference) { + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({}, kInvalidSnapshotId)); + ICEBERG_UNWRAP_OR_FAIL(auto metadata_table, + MetadataTable::Make(table)); + EXPECT_THAT(Scan(metadata_table, SnapshotSelection{.ref_name = "missing"}), + IsError(ErrorKind::kInvalidArgument)); +} + +TEST_P(PartitionsTableTest, ScansRefreshedSchemaAndPartitionSpec) { + auto first = MakeAppendSnapshotWithPartitionValues( + GetParam(), 10, std::nullopt, 1, + {{"old.parquet", PartitionValues(Literal::Int(7))}}, partitioned_spec_); + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({first}, 10, {}, partitioned_spec_)); + ICEBERG_UNWRAP_OR_FAIL(auto partitions, MetadataTable::Make(table)); + auto old_partitions_schema = partitions->schema(); + ICEBERG_UNWRAP_OR_FAIL(auto old_stream, Scan(partitions)); + auto old_reader = ::arrow::ImportRecordBatchReader(&old_stream); + ASSERT_TRUE(old_reader.ok()) << old_reader.status().ToString(); + + ASSERT_THAT(RefreshTableSchemaAndPartitionSpec(table, first), IsOk()); + + // The public builder factory must refresh the schema just like NewScan(). + ICEBERG_UNWRAP_OR_FAIL(auto builder, MetadataScanBuilder::Make(*partitions)); + ICEBERG_UNWRAP_OR_FAIL(auto scan, builder->Build()); + ICEBERG_UNWRAP_OR_FAIL(auto partitions_stream, scan->ToArrowArrayStream()); + ICEBERG_UNWRAP_OR_FAIL(auto partition_batches, + ReadAllBatches(std::move(partitions_stream))); + ASSERT_EQ(partition_batches.size(), 1); + // Files that only differ in the newly added partition field must not merge. + ASSERT_EQ(partition_batches[0]->num_rows(), 2); + auto partition = std::static_pointer_cast<::arrow::StructArray>( + partition_batches[0]->GetColumnByName("partition")); + ASSERT_EQ(partition->num_fields(), 2); + auto partition_ids = std::static_pointer_cast<::arrow::Int32Array>(partition->field(1)); + EXPECT_EQ(partition_ids->Value(0), 1); + EXPECT_EQ(partition_ids->Value(1), 2); + auto counts = std::static_pointer_cast<::arrow::Int32Array>( + partition_batches[0]->GetColumnByName("file_count")); + EXPECT_EQ(counts->Value(0), 1); + EXPECT_EQ(counts->Value(1), 1); + EXPECT_NE(partitions->schema(), old_partitions_schema); + + // Streams created before refresh retain their original metadata and schema. + auto batches = (*old_reader)->ToRecordBatches(); + ASSERT_TRUE(batches.ok()) << batches.status().ToString(); + ASSERT_EQ(batches->size(), 1); + EXPECT_EQ((*batches)[0]->num_rows(), 1); + auto old_partition = std::static_pointer_cast<::arrow::StructArray>( + (*batches)[0]->GetColumnByName("partition")); + EXPECT_EQ(old_partition->num_fields(), 1); +} + +TEST_P(PartitionsTableTest, ProjectsLiveFilesForPartitionAggregation) { + auto data_file = MakeDataFile("data.parquet", PartitionValues(Literal::Int(7)), + partitioned_spec_, 3); + data_file->column_sizes.emplace(1, 100); + data_file->value_counts.emplace(1, 3); + ICEBERG_UNWRAP_OR_FAIL(auto bound, Conversions::ToBytes(Literal::Int(42))); + data_file->lower_bounds.emplace(1, bound); + data_file->upper_bounds.emplace(1, bound); + auto data = WriteDataManifest(GetParam(), 10, + {MakeEntry(ManifestStatus::kAdded, 10, 1, data_file)}, + partitioned_spec_); + auto delete_file = MakeDataFile("delete.parquet", PartitionValues(Literal::Int(7)), + partitioned_spec_, 2); + delete_file->content = DataFile::Content::kPositionDeletes; + auto deletes = WriteDeleteManifest( + GetParam(), 10, {MakeEntry(ManifestStatus::kAdded, 10, 1, delete_file)}, + partitioned_spec_); + auto snapshot = std::make_shared(Snapshot{ + .snapshot_id = 10, + .sequence_number = 1, + .timestamp_ms = TimePointMsFromUnixMs(1000), + .manifest_list = WriteManifestList(GetParam(), 10, 0, 1, {data, deletes}), + }); + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({snapshot}, 10, {}, partitioned_spec_)); + ICEBERG_UNWRAP_OR_FAIL(auto all_files, + internal::LiveFiles(*table->metadata(), table->io(), snapshot)); + ICEBERG_UNWRAP_OR_FAIL(auto full_file, all_files->Next()); + ASSERT_TRUE(full_file.has_value()); + EXPECT_FALSE(full_file->file->column_sizes.empty()); + EXPECT_FALSE(full_file->file->lower_bounds.empty()); + + ICEBERG_UNWRAP_OR_FAIL(auto projected_files, + internal::LiveFiles(*table->metadata(), table->io(), snapshot, + {"content", "partition", "record_count", + "file_size_in_bytes"})); + ICEBERG_UNWRAP_OR_FAIL(auto projected_file, projected_files->Next()); + ASSERT_TRUE(projected_file.has_value()); + EXPECT_TRUE(projected_file->file->column_sizes.empty()); + EXPECT_TRUE(projected_file->file->value_counts.empty()); + EXPECT_TRUE(projected_file->file->lower_bounds.empty()); + EXPECT_TRUE(projected_file->file->upper_bounds.empty()); + EXPECT_EQ(projected_file->file->record_count, 3); + EXPECT_EQ(projected_file->snapshot_id, 10); + + ICEBERG_UNWRAP_OR_FAIL(auto partitions, MetadataTable::Make(table)); + ICEBERG_UNWRAP_OR_FAIL(auto stream, Scan(partitions)); + ICEBERG_UNWRAP_OR_FAIL(auto batches, ReadAllBatches(std::move(stream))); + ASSERT_EQ(batches.size(), 1); + ASSERT_EQ(batches[0]->num_rows(), 1); + for (const auto& [name, expected] : + std::vector>{{"record_count", 3}, + {"total_data_file_size_in_bytes", 10}, + {"position_delete_record_count", 2}, + {"last_updated_snapshot_id", 10}}) { + auto values = + std::static_pointer_cast<::arrow::Int64Array>(batches[0]->GetColumnByName(name)); + EXPECT_EQ(values->Value(0), expected) << name; + } +} + +TEST_P(PartitionsTableTest, GroupsNullPartitionValues) { + auto snapshot = MakeAppendSnapshotWithPartitionValues( + GetParam(), 10, std::nullopt, 1, + {{"s3://bucket/first.parquet", PartitionValues(Literal::Null(int32()))}, + {"s3://bucket/second.parquet", PartitionValues(Literal::Null(int32()))}}, + partitioned_spec_); + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({snapshot}, 10, {}, partitioned_spec_)); + ICEBERG_UNWRAP_OR_FAIL(auto partitions, MetadataTable::Make(table)); + + ICEBERG_UNWRAP_OR_FAIL(auto stream, Scan(partitions)); + ICEBERG_UNWRAP_OR_FAIL(auto batches, ReadAllBatches(std::move(stream))); + ASSERT_EQ(batches.size(), 1); + ASSERT_EQ(batches[0]->num_rows(), 1); + auto file_counts = std::static_pointer_cast<::arrow::Int32Array>( + batches[0]->GetColumnByName("file_count")); + EXPECT_EQ(file_counts->Value(0), 2); +} + +TEST_P(PartitionsTableTest, GroupsAllNaNsButDistinguishesSignedZeros) { + // Exercise both floating-point types, including NaNs with different payloads. + for (const bool use_float : {false, true}) { + schema_ = std::make_shared(std::vector{ + SchemaField::MakeRequired(1, "id", int32()), + SchemaField::MakeRequired( + 2, "data", use_float ? std::shared_ptr(float32()) : float64())}); + ICEBERG_UNWRAP_OR_FAIL( + partitioned_spec_, + PartitionSpec::Make(1, {PartitionField(2, 1000, "data", Transform::Identity())})); + const double nan = std::numeric_limits::quiet_NaN(); + auto literal = [use_float](double value) { + return use_float ? Literal::Float(static_cast(value)) + : Literal::Double(value); + }; + auto payload_nan = + use_float ? Literal::Float(std::bit_cast(uint32_t{0x7fc00001})) + : Literal::Double(std::bit_cast(uint64_t{0x7ff8000000000001})); + auto snapshot = MakeAppendSnapshotWithPartitionValues( + GetParam(), 10, std::nullopt, 1, + {{"positive_nan.parquet", PartitionValues(literal(nan))}, + {"negative_nan.parquet", PartitionValues(literal(-nan))}, + {"payload_nan.parquet", PartitionValues(payload_nan)}, + {"positive_zero.parquet", PartitionValues(literal(0.0))}, + {"negative_zero.parquet", PartitionValues(literal(-0.0))}}, + partitioned_spec_); + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({snapshot}, 10, {}, partitioned_spec_)); + ICEBERG_UNWRAP_OR_FAIL(auto partitions, MetadataTable::Make(table)); + ICEBERG_UNWRAP_OR_FAIL(auto stream, Scan(partitions)); + ICEBERG_UNWRAP_OR_FAIL(auto batches, ReadAllBatches(std::move(stream))); + ASSERT_EQ(batches.size(), 1); + ASSERT_EQ(batches[0]->num_rows(), 3); + auto counts = std::static_pointer_cast<::arrow::Int32Array>( + batches[0]->GetColumnByName("file_count")); + EXPECT_THAT( + (std::vector{counts->Value(0), counts->Value(1), counts->Value(2)}), + ::testing::UnorderedElementsAre(3, 1, 1)); + } +} + +TEST_P(PartitionsTableTest, ScansDecimalPartitions) { + const int128_t large = (static_cast(1) << 80) + 123; + const std::vector values{123, -456, large, -large, 0}; + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeDecimalPartitionTable(values)); + ICEBERG_UNWRAP_OR_FAIL(auto partitions, MetadataTable::Make(table)); + ICEBERG_UNWRAP_OR_FAIL(auto partition_stream, Scan(partitions)); + ICEBERG_UNWRAP_OR_FAIL(auto partition_batches, + ReadAllBatches(std::move(partition_stream))); + ASSERT_EQ(partition_batches.size(), 1); + ASSERT_EQ(partition_batches[0]->num_rows(), values.size()); + auto partition_values = std::static_pointer_cast<::arrow::StructArray>( + partition_batches[0]->GetColumnByName("partition")); + auto partition_amounts = + std::static_pointer_cast<::arrow::Decimal128Array>(partition_values->field(0)); + for (size_t index = 0; index < values.size(); ++index) { + ICEBERG_UNWRAP_OR_FAIL(auto expected, Decimal(values[index]).ToString(2)); + EXPECT_EQ(partition_amounts->FormatValue(index), expected); + } +} + +TEST_P(PartitionsTableTest, ScanEmpty) { + ICEBERG_UNWRAP_OR_FAIL(auto table, MakeTable({}, kInvalidSnapshotId)); + ICEBERG_UNWRAP_OR_FAIL(auto metadata_table, + MetadataTable::Make(table)); + ICEBERG_UNWRAP_OR_FAIL(auto stream, Scan(metadata_table)); + ICEBERG_UNWRAP_OR_FAIL(auto batches, ReadAllBatches(std::move(stream))); + EXPECT_TRUE(batches.empty()); +} + +TEST_P(PartitionsTableTest, TasksRetainProjectionAndMetadataWithCustomBatchSize) { + CheckTaskProjectionAndLifetime(MetadataTableKind::kPartitions); +} + +TEST_P(PartitionsTableTest, ConcurrentBuildersRetainTheirOwnSchemas) { + CheckConcurrentBuilders(MetadataTableKind::kPartitions); +} + +INSTANTIATE_TEST_SUITE_P(FormatVersions, PartitionsTableTest, ::testing::Values(2, 3)); + +} // namespace +} // namespace iceberg diff --git a/src/iceberg/type_fwd.h b/src/iceberg/type_fwd.h index 8e441d874..a44ebd282 100644 --- a/src/iceberg/type_fwd.h +++ b/src/iceberg/type_fwd.h @@ -291,6 +291,7 @@ class MetadataTable; class MetadataScan; class MetadataScanTask; class MetadataScanTaskReader; +class PartitionsTable; class RefsTable; class SnapshotsTable;