diff --git a/src/iceberg/CMakeLists.txt b/src/iceberg/CMakeLists.txt index 4cde4fb23..6bab2279e 100644 --- a/src/iceberg/CMakeLists.txt +++ b/src/iceberg/CMakeLists.txt @@ -63,6 +63,7 @@ set(ICEBERG_SOURCES inspect/metadata_table.cc inspect/metadata_scan_task.cc inspect/metadata_scan_task_reader.cc + inspect/metadata_table_util_internal.cc inspect/refs_table.cc inspect/snapshots_table.cc inheritable_metadata.cc diff --git a/src/iceberg/inspect/history_table.cc b/src/iceberg/inspect/history_table.cc index a6f616fdb..a6ffebbd8 100644 --- a/src/iceberg/inspect/history_table.cc +++ b/src/iceberg/inspect/history_table.cc @@ -77,7 +77,7 @@ HistoryTable::HistoryTable(std::shared_ptr table) HistoryTable::~HistoryTable() = default; -const std::shared_ptr& HistoryTable::schema() const { return CanonicalSchema(); } +std::shared_ptr HistoryTable::schema() const { return CanonicalSchema(); } const std::shared_ptr& HistoryTable::CanonicalSchema() { static const auto schema = std::make_shared(std::vector{ diff --git a/src/iceberg/inspect/history_table.h b/src/iceberg/inspect/history_table.h index ec18e36a4..f1b949bf9 100644 --- a/src/iceberg/inspect/history_table.h +++ b/src/iceberg/inspect/history_table.h @@ -46,7 +46,7 @@ class ICEBERG_EXPORT HistoryTable : public MetadataTable { MetadataTableKind kind() const noexcept override { return MetadataTableKind::kHistory; } - const std::shared_ptr& schema() const override; + std::shared_ptr schema() const override; protected: Result> MakeScan( diff --git a/src/iceberg/inspect/metadata_scan.cc b/src/iceberg/inspect/metadata_scan.cc index 2957a2ba4..03c03f567 100644 --- a/src/iceberg/inspect/metadata_scan.cc +++ b/src/iceberg/inspect/metadata_scan.cc @@ -107,20 +107,26 @@ class TaskArrowStreamSource { } // namespace -MetadataScanBuilder::MetadataScanBuilder(const MetadataTable& table) - : table_(&table), - context_{.kind = table.kind(), - .table_metadata = table.source_table()->metadata(), - .schema = table.schema(), - .projected_schema = table.schema()} {} +MetadataScanBuilder::MetadataScanBuilder(const MetadataTable& table, + internal::MetadataScanContext context) + : table_(&table), context_(std::move(context)) {} MetadataScanBuilder::~MetadataScanBuilder() = default; Result> MetadataScanBuilder::Make( const MetadataTable& table) { ICEBERG_PRECHECK(table.source_table() != nullptr, "Source table cannot be null"); - ICEBERG_PRECHECK(table.schema() != nullptr, "Metadata table schema cannot be null"); - return std::unique_ptr(new MetadataScanBuilder(table)); + auto metadata = table.source_table()->metadata(); + ICEBERG_PRECHECK(metadata != nullptr, "Table metadata cannot be null"); + ICEBERG_ASSIGN_OR_RAISE(auto schema, table.SchemaForScan(*metadata)); + ICEBERG_PRECHECK(schema != nullptr, "Metadata table schema cannot be null"); + internal::MetadataScanContext context{.kind = table.kind(), + .table_metadata = std::move(metadata), + .schema = schema, + .projected_schema = std::move(schema), + .io = table.source_table()->io()}; + return std::unique_ptr( + new MetadataScanBuilder(table, std::move(context))); } MetadataScanBuilder& MetadataScanBuilder::Select( diff --git a/src/iceberg/inspect/metadata_scan.h b/src/iceberg/inspect/metadata_scan.h index 5df4cc034..9289e6515 100644 --- a/src/iceberg/inspect/metadata_scan.h +++ b/src/iceberg/inspect/metadata_scan.h @@ -52,6 +52,7 @@ struct MetadataScanContext { std::optional snapshot_id; std::string ref_name; std::optional as_of_time_in_millis; + std::shared_ptr io; }; } // namespace internal @@ -97,7 +98,7 @@ class ICEBERG_EXPORT MetadataScanBuilder : public ErrorCollector { Result> Build(); private: - explicit MetadataScanBuilder(const MetadataTable& table); + MetadataScanBuilder(const MetadataTable& table, internal::MetadataScanContext context); const MetadataTable* table_; internal::MetadataScanContext context_; diff --git a/src/iceberg/inspect/metadata_scan_task.h b/src/iceberg/inspect/metadata_scan_task.h index 6e0dac54e..fc746efcd 100644 --- a/src/iceberg/inspect/metadata_scan_task.h +++ b/src/iceberg/inspect/metadata_scan_task.h @@ -37,6 +37,8 @@ class ICEBERG_EXPORT MetadataScanTask { enum class TaskKind { /// StaticMetadataTask reads in-memory table metadata (snapshots, history, refs). kStatic, + /// SnapshotMetadataTask reads snapshot manifests and their file entries. + kSnapshot, }; virtual ~MetadataScanTask() = default; diff --git a/src/iceberg/inspect/metadata_scan_task_reader.cc b/src/iceberg/inspect/metadata_scan_task_reader.cc index 74acaa8f3..abe69c498 100644 --- a/src/iceberg/inspect/metadata_scan_task_reader.cc +++ b/src/iceberg/inspect/metadata_scan_task_reader.cc @@ -25,6 +25,7 @@ #include "iceberg/inspect/history_table.h" #include "iceberg/inspect/metadata_scan_task.h" #include "iceberg/inspect/metadata_table.h" +#include "iceberg/inspect/metadata_table_util_internal.h" #include "iceberg/inspect/refs_table.h" #include "iceberg/inspect/snapshots_table.h" #include "iceberg/util/checked_cast.h" @@ -47,6 +48,11 @@ Result> MetadataScanTaskReader::Make( Result MetadataScanTaskReader::Open( const MetadataScanTask& task) const { switch (task.task_kind()) { + case MetadataScanTask::TaskKind::kSnapshot: { + const auto& snapshot_task = + internal::checked_cast(task); + return snapshot_task.Open(options_.batch_size); + } case MetadataScanTask::TaskKind::kStatic: { const auto& static_task = internal::checked_cast(task); switch (static_task.table_kind()) { diff --git a/src/iceberg/inspect/metadata_scan_util_internal.h b/src/iceberg/inspect/metadata_scan_util_internal.h index 5cd2058f9..65e1f3ef3 100644 --- a/src/iceberg/inspect/metadata_scan_util_internal.h +++ b/src/iceberg/inspect/metadata_scan_util_internal.h @@ -40,6 +40,7 @@ #include "iceberg/schema_internal.h" #include "iceberg/type.h" #include "iceberg/util/macros.h" +#include "iceberg/util/stream.h" namespace iceberg::internal { @@ -181,4 +182,79 @@ Result MakeMetadataBatchStream( batch_size, std::forward(get_row))); } +/// Stream manifest-backed rows with the same batch-size and ownership contract. +template +class MetadataRowsSource : public Stream { + public: + MetadataRowsSource(std::shared_ptr schema, + std::unique_ptr> rows, int64_t batch_size, + AppendRow append_row) + : schema_(std::move(schema)), + rows_(std::move(rows)), + batch_size_(batch_size), + append_row_(std::move(append_row)) {} + + Status Close() { + rows_.reset(); + return {}; + } + + Result Schema() { + ArrowSchema schema{}; + ICEBERG_RETURN_UNEXPECTED(ToArrowSchema(*schema_, &schema)); + return schema; + } + + protected: + Result> NextImpl() override { + ICEBERG_PRECHECK(rows_ != nullptr, "Cannot read from a closed metadata stream"); + ICEBERG_ASSIGN_OR_RAISE(auto row, rows_->Next()); + if (!row) { + return std::nullopt; + } + ICEBERG_ASSIGN_OR_RAISE(auto builder, ArrowRowBuilder::Make(*schema_)); + ICEBERG_RETURN_UNEXPECTED(append_row_(builder, *row)); + while (builder.num_rows() < batch_size_) { + ICEBERG_ASSIGN_OR_RAISE(auto next_row, rows_->Next()); + if (!next_row) { + break; + } + ICEBERG_RETURN_UNEXPECTED(append_row_(builder, *next_row)); + } + ICEBERG_ASSIGN_OR_RAISE(auto batch, std::move(builder).Finish()); + return batch; + } + + private: + const std::shared_ptr schema_; + std::unique_ptr> rows_; + const int64_t batch_size_; + AppendRow append_row_; +}; + +template +Result MakeMetadataBatchStream(const std::shared_ptr& schema, + std::unique_ptr> rows, + int64_t batch_size, + AppendRow append_row) { + ICEBERG_PRECHECK(batch_size > 0, "Batch size must be positive"); + ICEBERG_PRECHECK(schema != nullptr, "Metadata stream schema cannot be null"); + ICEBERG_PRECHECK(rows != nullptr, "Metadata row source cannot be null"); + if constexpr (requires { static_cast(append_row); }) { + ICEBERG_PRECHECK(static_cast(append_row), + "Metadata row callback cannot be empty"); + } + return MakeArrowArrayStream(std::make_unique>( + schema, std::move(rows), batch_size, std::move(append_row))); +} + +template +Result MakeMetadataBatchStream(const std::shared_ptr& schema, + std::vector rows, + int64_t batch_size, + AppendRow append_row) { + return MakeMetadataBatchStream(schema, Stream::FromVector(std::move(rows)), + batch_size, std::move(append_row)); +} + } // namespace iceberg::internal diff --git a/src/iceberg/inspect/metadata_table.cc b/src/iceberg/inspect/metadata_table.cc index 012ffaa99..57d8cdad3 100644 --- a/src/iceberg/inspect/metadata_table.cc +++ b/src/iceberg/inspect/metadata_table.cc @@ -53,6 +53,11 @@ const std::shared_ptr& MetadataTable::source_table() const { return source_table_; } +Result> MetadataTable::SchemaForScan( + const TableMetadata& metadata) const { + return schema(); +} + Result> MetadataTable::NewScan() const { return MetadataScanBuilder::Make(*this); } diff --git a/src/iceberg/inspect/metadata_table.h b/src/iceberg/inspect/metadata_table.h index 83f44f5c1..99e98ef1e 100644 --- a/src/iceberg/inspect/metadata_table.h +++ b/src/iceberg/inspect/metadata_table.h @@ -78,8 +78,8 @@ class ICEBERG_EXPORT MetadataTable { /// \brief Return this metadata table's kind. virtual Kind kind() const noexcept = 0; - /// \brief Return the complete canonical schema of this metadata table. - virtual const std::shared_ptr& schema() const = 0; + /// \brief Return shared ownership of this metadata table's canonical schema. + virtual std::shared_ptr schema() const = 0; /// \brief Return the source table whose metadata is exposed. const std::shared_ptr& source_table() const; @@ -93,6 +93,10 @@ class ICEBERG_EXPORT MetadataTable { protected: explicit MetadataTable(std::shared_ptr source_table); + /// \brief Derive a scan schema from the captured source metadata. + virtual Result> SchemaForScan( + const TableMetadata& metadata) const; + /// \brief Plan a scan using the captured metadata, projection, and snapshot selection. virtual Result> MakeScan( internal::MetadataScanContext context) const = 0; diff --git a/src/iceberg/inspect/metadata_table_util_internal.cc b/src/iceberg/inspect/metadata_table_util_internal.cc new file mode 100644 index 000000000..a5e00dcab --- /dev/null +++ b/src/iceberg/inspect/metadata_table_util_internal.cc @@ -0,0 +1,397 @@ +/* + * 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/metadata_table_util_internal.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "iceberg/arrow_row_builder_internal.h" +#include "iceberg/constants.h" +#include "iceberg/file_format.h" +#include "iceberg/manifest/manifest_list.h" +#include "iceberg/manifest/manifest_reader.h" +#include "iceberg/nanoarrow_status_internal.h" +#include "iceberg/partition_spec.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/transform.h" +#include "iceberg/type.h" +#include "iceberg/util/checked_cast.h" +#include "iceberg/util/conversions.h" +#include "iceberg/util/macros.h" +#include "iceberg/util/snapshot_util.h" + +namespace iceberg::internal { +namespace { + +Result> SnapshotAtRef(const TableMetadata& metadata, + std::string_view ref_name) { + if (ref_name.empty() || ref_name == SnapshotRef::kMainBranch) { + if (metadata.current_snapshot_id == kInvalidSnapshotId) { + return std::shared_ptr{nullptr}; + } + return metadata.SnapshotById(metadata.current_snapshot_id); + } + + auto ref = metadata.refs.find(std::string(ref_name)); + ICEBERG_PRECHECK(ref != metadata.refs.end(), "Cannot find snapshot reference '{}'", + ref_name); + ICEBERG_PRECHECK(ref->second != nullptr, "Snapshot reference '{}' is null", ref_name); + return metadata.SnapshotById(ref->second->snapshot_id); +} + +} // namespace + +Status AppendLiteral(ArrowArray* array, const Literal& literal) { + if (literal.IsNull()) { + return AppendNull(array); + } + if (literal.IsAboveMax() || literal.IsBelowMin()) { + return InvalidArgument("Cannot append non-value partition literal {}", + literal.ToString()); + } + + switch (literal.type()->type_id()) { + case TypeId::kBoolean: + return AppendBoolean(array, std::get(literal.value())); + case TypeId::kInt: + case TypeId::kDate: + return AppendInt(array, std::get(literal.value())); + case TypeId::kLong: + case TypeId::kTime: + case TypeId::kTimestamp: + case TypeId::kTimestampTz: + case TypeId::kTimestampNs: + case TypeId::kTimestampTzNs: + return AppendInt(array, std::get(literal.value())); + case TypeId::kFloat: + return AppendDouble(array, std::get(literal.value())); + case TypeId::kDouble: + return AppendDouble(array, std::get(literal.value())); + case TypeId::kString: + return AppendString(array, std::get(literal.value())); + case TypeId::kBinary: + case TypeId::kFixed: + return AppendBytes(array, std::get>(literal.value())); + case TypeId::kDecimal: + return AppendDecimal(array, std::get(literal.value())); + case TypeId::kUuid: + return AppendBytes(array, std::get(literal.value()).bytes()); + case TypeId::kUnknown: + case TypeId::kStruct: + case TypeId::kList: + case TypeId::kMap: + case TypeId::kVariant: + case TypeId::kGeometry: + case TypeId::kGeography: + return NotSupported("Cannot append partition literal of type {}", + literal.type()->ToString()); + } + std::unreachable(); +} + +namespace { +class SnapshotMetadataScan : public MetadataScan { + public: + SnapshotMetadataScan(MetadataScanContext context, + SnapshotMetadataTask::ReadRowsFunction read_rows) + : MetadataScan(context), + task_(std::make_shared(std::move(context), read_rows)) {} + + Result PlanTasksStream() const override { + return MetadataScanTaskStream::FromVector({task_}); + } + + private: + const std::shared_ptr task_; +}; +} // namespace + +Result> MakeSnapshotMetadataScan( + MetadataScanContext context, SnapshotMetadataTask::ReadRowsFunction read_rows) { + ICEBERG_PRECHECK(read_rows != nullptr, "Snapshot row reader cannot be null"); + ICEBERG_PRECHECK(context.table_metadata != nullptr, "Table metadata cannot be null"); + ICEBERG_PRECHECK(context.schema != nullptr, "Metadata table schema cannot be null"); + ICEBERG_PRECHECK(context.projected_schema != nullptr, + "Projected schema cannot be null"); + ICEBERG_PRECHECK(context.io != nullptr, "FileIO cannot be null"); + for (const auto& projected : context.projected_schema->fields()) { + ICEBERG_ASSIGN_OR_RAISE(auto field, + context.schema->GetFieldById(projected.field_id())); + ICEBERG_PRECHECK(field.has_value(), "Unknown metadata column: {}", projected.name()); + if (*field->get().type() != *projected.type()) { + return NotSupported("Partial projection of metadata column is not supported: {}", + projected.name()); + } + ICEBERG_PRECHECK(!field->get().optional() || projected.optional(), + "Nullable metadata column cannot be projected as required: {}", + projected.name()); + } + ICEBERG_RETURN_UNEXPECTED( + ResolveMetadataTableSnapshot(*context.table_metadata, context)); + return std::make_unique(std::move(context), read_rows); +} + +Result> ResolveMetadataTableSnapshot( + const TableMetadata& metadata, const MetadataScanContext& selection) { + if (!selection.ref_name.empty() && selection.ref_name != SnapshotRef::kMainBranch) { + ICEBERG_PRECHECK( + !selection.snapshot_id.has_value() && !selection.as_of_time_in_millis.has_value(), + "Cannot combine a snapshot reference with a snapshot ID or timestamp"); + } + if (selection.snapshot_id.has_value()) { + return metadata.SnapshotById(*selection.snapshot_id); + } + if (selection.as_of_time_in_millis.has_value()) { + ICEBERG_ASSIGN_OR_RAISE( + auto snapshot_id, + SnapshotUtil::SnapshotIdAsOfTime( + metadata, TimePointMsFromUnixMs(*selection.as_of_time_in_millis))); + return metadata.SnapshotById(snapshot_id); + } + return SnapshotAtRef(metadata, selection.ref_name); +} + +Result> UnifiedPartitionType(const TableMetadata& metadata) { + ICEBERG_ASSIGN_OR_RAISE(auto schema, metadata.Schema()); + auto specs = metadata.partition_specs; + for (const auto& spec : specs) { + ICEBERG_PRECHECK(spec != nullptr, "Partition spec cannot be null"); + } + std::ranges::sort(specs, std::greater{}, &PartitionSpec::spec_id); + + std::unordered_set active_field_ids; + for (const auto& spec : specs) { + for (const auto& field : spec->fields()) { + ICEBERG_PRECHECK(field.transform() != nullptr, + "Partition field {} has a null transform", field.field_id()); + ICEBERG_CHECK(field.transform()->transform_type() != TransformType::kUnknown, + "Cannot build table partition type with unknown transform '{}'", + field.transform()->ToString()); + ICEBERG_ASSIGN_OR_RAISE(auto source_field, + schema->FindFieldById(field.source_id())); + if (source_field.has_value()) { + active_field_ids.insert(field.field_id()); + } + } + } + + struct ProjectedField { + const PartitionField* definition; + std::string name; + std::shared_ptr type; + }; + std::map fields_by_id; + for (const auto& spec : specs) { + ICEBERG_ASSIGN_OR_RAISE(auto spec_type, spec->PartitionType(*schema)); + ICEBERG_PRECHECK(spec_type->fields().size() == spec->fields().size(), + "Partition spec {} has mismatched field and type counts", + spec->spec_id()); + for (size_t index = 0; index < spec->fields().size(); ++index) { + const auto& partition_field = spec->fields()[index]; + if (!active_field_ids.contains(partition_field.field_id())) { + continue; + } + + const auto& spec_field = spec_type->fields()[index]; + auto [iter, inserted] = + fields_by_id.try_emplace(partition_field.field_id(), + ProjectedField{.definition = &partition_field, + .name = std::string(spec_field.name()), + .type = spec_field.type()}); + if (!inserted) { + const auto& existing = *iter->second.definition; + const auto current_transform = partition_field.transform()->transform_type(); + const auto existing_transform = existing.transform()->transform_type(); + const bool compatible_transform = + *partition_field.transform() == *existing.transform() || + current_transform == TransformType::kVoid || + existing_transform == TransformType::kVoid; + ICEBERG_CHECK( + partition_field.source_id() == existing.source_id() && compatible_transform, + "Conflicting partition fields with ID {}: '{}' and '{}'", + partition_field.field_id(), partition_field.ToString(), existing.ToString()); + + if (existing_transform == TransformType::kVoid && + current_transform != TransformType::kVoid) { + iter->second.definition = &partition_field; + iter->second.type = spec_field.type(); + } + } + } + } + + std::vector fields; + fields.reserve(fields_by_id.size()); + for (auto& [_, field] : fields_by_id) { + fields.emplace_back(field.definition->field_id(), std::move(field.name), + std::move(field.type), /*optional=*/true); + } + return std::make_shared(std::move(fields)); +} + +Result ProjectPartitionValues(const StructType& partition_type, + const PartitionSpec& spec, + const PartitionValues& values) { + ICEBERG_PRECHECK(values.num_fields() == spec.fields().size(), + "Partition has {} values but spec {} has {} fields", + values.num_fields(), spec.spec_id(), spec.fields().size()); + + std::unordered_map positions; + positions.reserve(spec.fields().size()); + for (size_t index = 0; index < spec.fields().size(); ++index) { + positions.emplace(spec.fields()[index].field_id(), index); + } + + std::vector projected; + projected.reserve(partition_type.fields().size()); + for (const auto& field : partition_type.fields()) { + auto target_type = checked_pointer_cast(field.type()); + auto position = positions.find(field.field_id()); + if (position == positions.end()) { + projected.push_back(Literal::Null(std::move(target_type))); + continue; + } + ICEBERG_ASSIGN_OR_RAISE(auto value, values.ValueAt(position->second)); + if (value.get().IsNull()) { + projected.push_back(Literal::Null(std::move(target_type))); + } else if (value.get().type()->type_id() == target_type->type_id()) { + projected.push_back(value.get()); + } else { + ICEBERG_ASSIGN_OR_RAISE(auto coerced, value.get().CastTo(target_type)); + projected.push_back(std::move(coerced)); + } + } + return PartitionValues(std::move(projected)); +} + +Status AppendPartitionValues(ArrowArray* array, const StructType& partition_type, + const PartitionValues& values) { + ICEBERG_PRECHECK(array != nullptr, "Partition Arrow array cannot be null"); + ICEBERG_PRECHECK(std::cmp_equal(array->n_children, partition_type.fields().size()), + "Partition Arrow array has {} fields but schema has {}", + array->n_children, partition_type.fields().size()); + ICEBERG_PRECHECK(values.num_fields() == partition_type.fields().size(), + "Partition has {} values but schema has {} fields", + values.num_fields(), partition_type.fields().size()); + + for (size_t index = 0; index < values.num_fields(); ++index) { + ICEBERG_ASSIGN_OR_RAISE(auto value, values.ValueAt(index)); + ICEBERG_RETURN_UNEXPECTED(AppendLiteral(array->children[index], value.get())); + } + ICEBERG_NANOARROW_RETURN_UNEXPECTED(ArrowArrayFinishElement(array)); + return {}; +} + +namespace { + +class LiveFilesIterator : public Stream { + public: + LiveFilesIterator(std::shared_ptr io, std::shared_ptr schema, + std::unordered_map> specs, + std::vector manifests, std::vector columns) + : io_(std::move(io)), + schema_(std::move(schema)), + specs_(std::move(specs)), + manifests_(std::move(manifests)), + columns_(std::move(columns)) {} + + protected: + Result> NextImpl() override { + while (entry_index_ == entries_.size()) { + entries_.clear(); + entry_index_ = 0; + if (manifest_index_ == manifests_.size()) { + return std::nullopt; + } + const auto& manifest = manifests_[manifest_index_++]; + auto spec = specs_.find(manifest.partition_spec_id); + ICEBERG_CHECK(spec != specs_.end(), + "Cannot find partition spec {} for manifest '{}'", + manifest.partition_spec_id, manifest.manifest_path); + ICEBERG_PRECHECK(spec->second != nullptr, "Partition spec {} is null", + manifest.partition_spec_id); + current_spec_ = spec->second; + ICEBERG_ASSIGN_OR_RAISE(auto reader, + ManifestReader::Make(manifest, io_, schema_, specs_)); + if (!columns_.empty()) { + reader->Select(columns_); + } + ICEBERG_ASSIGN_OR_RAISE(entries_, reader->LiveEntries()); + } + auto& entry = entries_[entry_index_++]; + ICEBERG_PRECHECK(entry.data_file != nullptr, + "Manifest contains an entry with no data file"); + return LiveFile{.file = std::move(entry.data_file), + .spec = current_spec_, + .snapshot_id = entry.snapshot_id}; + } + + private: + std::shared_ptr io_; + std::shared_ptr schema_; + std::unordered_map> specs_; + std::vector manifests_; + std::vector columns_; + std::vector entries_; + std::shared_ptr current_spec_; + size_t manifest_index_ = 0; + size_t entry_index_ = 0; +}; + +} // namespace + +Result>> LiveFiles( + const TableMetadata& metadata, const std::shared_ptr& io, + const std::shared_ptr& snapshot, std::vector columns) { + ICEBERG_ASSIGN_OR_RAISE(auto schema, metadata.Schema()); + std::unordered_map> specs; + for (const auto& spec : metadata.partition_specs) { + ICEBERG_PRECHECK(spec != nullptr, "Partition spec cannot be null"); + specs.emplace(spec->spec_id(), spec); + } + std::vector manifests; + if (snapshot != nullptr) { + SnapshotReader snapshot_reader(snapshot.get()); + ICEBERG_ASSIGN_OR_RAISE(auto snapshot_manifests, snapshot_reader.Manifests(io)); + manifests.assign(snapshot_manifests.begin(), snapshot_manifests.end()); + } + return std::make_unique(io, std::move(schema), std::move(specs), + std::move(manifests), std::move(columns)); +} + +} // namespace iceberg::internal diff --git a/src/iceberg/inspect/metadata_table_util_internal.h b/src/iceberg/inspect/metadata_table_util_internal.h new file mode 100644 index 000000000..b06cf8031 --- /dev/null +++ b/src/iceberg/inspect/metadata_table_util_internal.h @@ -0,0 +1,104 @@ +/* + * 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 "iceberg/iceberg_export.h" +#include "iceberg/inspect/metadata_scan.h" +#include "iceberg/inspect/metadata_scan_task.h" +#include "iceberg/manifest/manifest_entry.h" +#include "iceberg/result.h" +#include "iceberg/row/partition_values.h" +#include "iceberg/type_fwd.h" +#include "iceberg/util/stream.h" + +namespace iceberg { +class ArrowRowBuilder; +} + +namespace iceberg::internal { + +/// A task retaining the captured metadata, projection, and manifest IO. +class SnapshotMetadataTask : public MetadataScanTask { + public: + using ReadRowsFunction = Result (*)(const MetadataScanContext&, + const std::shared_ptr&, + int64_t); + + SnapshotMetadataTask(MetadataScanContext context, ReadRowsFunction read_rows) + : context(std::move(context)), read_rows_(read_rows) {} + + Result Open(int64_t batch_size) const { + return read_rows_(context, context.io, batch_size); + } + + TaskKind task_kind() const override { return TaskKind::kSnapshot; } + MetadataTableKind table_kind() const override { return context.kind; } + + const MetadataScanContext context; + + private: + ReadRowsFunction read_rows_; +}; + +ICEBERG_EXPORT Result> MakeSnapshotMetadataScan( + MetadataScanContext context, SnapshotMetadataTask::ReadRowsFunction read_rows); + +struct LiveFile { + std::shared_ptr file; + std::shared_ptr spec; + std::optional snapshot_id; +}; + +/// \brief Resolve a time-travel selection to a snapshot. +ICEBERG_EXPORT Result> ResolveMetadataTableSnapshot( + const TableMetadata& metadata, const MetadataScanContext& selection); + +/// \brief Build the Java-compatible union of active partition fields across all specs. +ICEBERG_EXPORT Result> UnifiedPartitionType( + const TableMetadata& metadata); + +/// \brief Project values written with one spec into the table-wide partition type. +ICEBERG_EXPORT Result ProjectPartitionValues( + const StructType& partition_type, const PartitionSpec& spec, + const PartitionValues& values); + +/// \brief Append a primitive partition value or decoded metric bound. +ICEBERG_EXPORT Status AppendLiteral(ArrowArray* array, const Literal& literal); + +/// \brief Append partition values to an Arrow struct builder. +ICEBERG_EXPORT Status AppendPartitionValues(ArrowArray* array, + const StructType& partition_type, + const PartitionValues& values); + +/// \brief Iterate over live files, reading one manifest at a time. +/// An empty column selection reads all data-file fields. Each manifest's entries +/// are materialized by ManifestReader, independently of the Arrow batch size. +ICEBERG_EXPORT Result>> LiveFiles( + const TableMetadata& metadata, const std::shared_ptr& io, + const std::shared_ptr& snapshot, std::vector columns = {}); + +} // namespace iceberg::internal diff --git a/src/iceberg/inspect/refs_table.cc b/src/iceberg/inspect/refs_table.cc index c57e092bd..25d064da3 100644 --- a/src/iceberg/inspect/refs_table.cc +++ b/src/iceberg/inspect/refs_table.cc @@ -91,7 +91,7 @@ RefsTable::RefsTable(std::shared_ptr table) RefsTable::~RefsTable() = default; -const std::shared_ptr& RefsTable::schema() const { return CanonicalSchema(); } +std::shared_ptr RefsTable::schema() const { return CanonicalSchema(); } const std::shared_ptr& RefsTable::CanonicalSchema() { static const auto schema = std::make_shared(std::vector{ diff --git a/src/iceberg/inspect/refs_table.h b/src/iceberg/inspect/refs_table.h index baa11f192..2628ffe6b 100644 --- a/src/iceberg/inspect/refs_table.h +++ b/src/iceberg/inspect/refs_table.h @@ -46,7 +46,7 @@ class ICEBERG_EXPORT RefsTable : public MetadataTable { MetadataTableKind kind() const noexcept override { return MetadataTableKind::kRefs; } - const std::shared_ptr& schema() const override; + std::shared_ptr schema() const override; protected: Result> MakeScan( diff --git a/src/iceberg/inspect/snapshots_table.cc b/src/iceberg/inspect/snapshots_table.cc index 7eb204b43..c54a71778 100644 --- a/src/iceberg/inspect/snapshots_table.cc +++ b/src/iceberg/inspect/snapshots_table.cc @@ -82,9 +82,7 @@ SnapshotsTable::SnapshotsTable(std::shared_ptr table) SnapshotsTable::~SnapshotsTable() = default; -const std::shared_ptr& SnapshotsTable::schema() const { - return CanonicalSchema(); -} +std::shared_ptr SnapshotsTable::schema() const { return CanonicalSchema(); } const std::shared_ptr& SnapshotsTable::CanonicalSchema() { static const auto schema = std::make_shared(std::vector{ diff --git a/src/iceberg/inspect/snapshots_table.h b/src/iceberg/inspect/snapshots_table.h index 4bfaa554f..2293ad929 100644 --- a/src/iceberg/inspect/snapshots_table.h +++ b/src/iceberg/inspect/snapshots_table.h @@ -48,7 +48,7 @@ class ICEBERG_EXPORT SnapshotsTable : public MetadataTable { return MetadataTableKind::kSnapshots; } - const std::shared_ptr& schema() const override; + std::shared_ptr schema() const override; protected: Result> MakeScan( diff --git a/src/iceberg/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 1874fc25e..01e3b505d 100644 --- a/src/iceberg/test/CMakeLists.txt +++ b/src/iceberg/test/CMakeLists.txt @@ -206,6 +206,7 @@ if(ICEBERG_BUILD_BUNDLE) add_iceberg_test(metadata_table_test USE_BUNDLE SOURCES + inspect/metadata_table_util_test.cc inspect/history_table_test.cc inspect/metadata_table_test.cc inspect/refs_table_test.cc 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/metadata_table_util_test.cc b/src/iceberg/test/inspect/metadata_table_util_test.cc new file mode 100644 index 000000000..06fcdb08f --- /dev/null +++ b/src/iceberg/test/inspect/metadata_table_util_test.cc @@ -0,0 +1,219 @@ +/* + * 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 + +#include "iceberg/inspect/metadata_scan_task_reader.h" +#include "iceberg/inspect/metadata_scan_util_internal.h" +#include "iceberg/inspect/metadata_table.h" +#include "iceberg/inspect/metadata_table_util_internal.h" +#include "iceberg/partition_spec.h" +#include "iceberg/schema.h" +#include "iceberg/schema_field.h" +#include "iceberg/test/inspect/metadata_table_test_base.h" +#include "iceberg/test/matchers.h" +#include "iceberg/transform.h" +#include "iceberg/type.h" + +namespace iceberg::internal { +namespace { + +Status AppendIntRow(ArrowRowBuilder& builder, const int64_t& row) { + ICEBERG_RETURN_UNEXPECTED(AppendInt(builder.column(0), row)); + return builder.FinishRow(); +} + +std::shared_ptr IntSchema() { + return std::make_shared( + std::vector{SchemaField::MakeRequired(1, "id", int64())}); +} + +class CountingRows : public Stream { + public: + explicit CountingRows(std::shared_ptr reads) : reads_(std::move(reads)) {} + + protected: + Result> NextImpl() override { + ++*reads_; + if (next_ == 5) { + return std::nullopt; + } + return next_++; + } + + private: + std::shared_ptr reads_; + int64_t next_ = 0; +}; + +TEST(MetadataRowsTest, RejectsInvalidInputsBeforeExport) { + EXPECT_THAT(MakeMetadataBatchStream(IntSchema(), std::unique_ptr>{}, 2, + AppendIntRow), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT(MakeMetadataBatchStream(std::shared_ptr{}, std::vector{1}, + 2, AppendIntRow), + IsError(ErrorKind::kInvalidArgument)); + EXPECT_THAT( + MakeMetadataBatchStream(IntSchema(), std::vector{1}, 0, AppendIntRow), + IsError(ErrorKind::kInvalidArgument)); + decltype(&AppendIntRow) callback = nullptr; + EXPECT_THAT(MakeMetadataBatchStream(IntSchema(), std::vector{1}, 2, callback), + IsError(ErrorKind::kInvalidArgument)); + std::function empty; + EXPECT_THAT(MakeMetadataBatchStream(IntSchema(), std::vector{1}, 2, empty), + IsError(ErrorKind::kInvalidArgument)); +} + +TEST(MetadataRowsTest, LazyBatchesAndExhaustion) { + auto reads = std::make_shared(0); + std::unique_ptr> rows = std::make_unique(reads); + ICEBERG_UNWRAP_OR_FAIL(auto stream, MakeMetadataBatchStream( + IntSchema(), std::move(rows), 2, AppendIntRow)); + nanoarrow::UniqueArrayStream owned(&stream); + nanoarrow::UniqueArray batch; + EXPECT_EQ(*reads, 0); + for (int length : {2, 2, 1}) { + ASSERT_EQ(owned->get_next(owned.get(), batch.get()), 0); + ASSERT_NE(batch->release, nullptr); + EXPECT_EQ(batch->length, length); + batch.reset(); + } + for (int i = 0; i < 3; ++i) { + EXPECT_EQ(owned->get_next(owned.get(), batch.get()), 0); + EXPECT_EQ(batch->release, nullptr); + } + EXPECT_EQ(*reads, 6); +} + +TEST(MetadataRowsTest, AppendFailureIsTerminal) { + auto reads = std::make_shared(0); + std::unique_ptr> rows = std::make_unique(reads); + auto append = [](ArrowRowBuilder& builder, const int64_t& row) -> Status { + if (row == 3) { + return InvalidArrowData("Row append failed"); + } + return AppendIntRow(builder, row); + }; + ICEBERG_UNWRAP_OR_FAIL( + auto stream, MakeMetadataBatchStream(IntSchema(), std::move(rows), 2, append)); + nanoarrow::UniqueArrayStream owned(&stream); + nanoarrow::UniqueArray batch; + ASSERT_EQ(owned->get_next(owned.get(), batch.get()), 0); + batch.reset(); + for (int i = 0; i < 3; ++i) { + EXPECT_EQ(owned->get_next(owned.get(), batch.get()), EIO); + EXPECT_EQ(batch->release, nullptr); + EXPECT_STREQ(owned->get_last_error(owned.get()), "Row append failed"); + } + EXPECT_EQ(*reads, 4); +} + +Result ReadIntTaskRows(const MetadataScanContext& context, + const std::shared_ptr& io, + int64_t batch_size) { + ICEBERG_PRECHECK(io != nullptr && context.table_metadata != nullptr, + "Task lost its input owners"); + return MakeMetadataBatchStream(context.projected_schema, std::vector{1, 2, 3}, + batch_size, AppendIntRow); +} + +class SnapshotMetadataSupportTest : public MetadataTableTestBase { + protected: + MetadataScanContext Context() { + auto schema = IntSchema(); + return {.kind = MetadataTableKind::kFiles, + .table_metadata = metadata_, + .schema = schema, + .projected_schema = schema, + .io = io_}; + } +}; + +TEST_F(SnapshotMetadataSupportTest, TasksRetainOwnersAndCanBeReopened) { + ICEBERG_UNWRAP_OR_FAIL(auto scan, MakeSnapshotMetadataScan(Context(), ReadIntTaskRows)); + ICEBERG_UNWRAP_OR_FAIL(auto tasks, scan->PlanTasks()); + ASSERT_EQ(tasks.size(), 1); + scan.reset(); + table_.reset(); + metadata_.reset(); + io_.reset(); + for (int attempt = 0; attempt < 2; ++attempt) { + ICEBERG_UNWRAP_OR_FAIL(auto reader, MetadataScanTaskReader::Make({.batch_size = 2})); + ICEBERG_UNWRAP_OR_FAIL(auto stream, reader->Open(*tasks[0])); + reader.reset(); + ICEBERG_UNWRAP_OR_FAIL(auto batches, ReadAllBatches(std::move(stream))); + ASSERT_EQ(batches.size(), 2); + EXPECT_EQ(batches[0]->num_rows(), 2); + EXPECT_EQ(batches[1]->num_rows(), 1); + ASSERT_EQ(batches[0]->column(0)->type_id(), ::arrow::Type::INT64); + auto ids = std::static_pointer_cast<::arrow::Int64Array>(batches[0]->column(0)); + EXPECT_EQ(ids->Value(0), 1); + } +} + +TEST_F(SnapshotMetadataSupportTest, RejectsInvalidProjectionAndReader) { + EXPECT_THAT(MakeSnapshotMetadataScan(Context(), nullptr), + IsError(ErrorKind::kInvalidArgument)); + auto context = Context(); + context.projected_schema = std::make_shared( + std::vector{SchemaField::MakeRequired(99, "missing", int64())}); + EXPECT_THAT(MakeSnapshotMetadataScan(context, ReadIntTaskRows), + IsError(ErrorKind::kInvalidArgument)); + context = Context(); + context.ref_name = "dev"; + context.snapshot_id = 1; + EXPECT_THAT(MakeSnapshotMetadataScan(context, ReadIntTaskRows), + IsError(ErrorKind::kInvalidArgument)); +} + +TEST_F(SnapshotMetadataSupportTest, ProjectsAcrossPartitionSpecEvolution) { + ICEBERG_UNWRAP_OR_FAIL( + auto old_spec, + PartitionSpec::Make(1, {PartitionField(1, 1000, "old_id", Transform::Identity())})); + ICEBERG_UNWRAP_OR_FAIL( + auto new_spec, + PartitionSpec::Make(2, {PartitionField(1, 1000, "new_id", Transform::Identity()), + PartitionField(2, 1001, "data", Transform::Identity())})); + metadata_->partition_specs.clear(); + metadata_->partition_specs.push_back(std::move(old_spec)); + metadata_->partition_specs.push_back(std::move(new_spec)); + ICEBERG_UNWRAP_OR_FAIL(auto type, UnifiedPartitionType(*metadata_)); + ASSERT_EQ(type->fields().size(), 2); + EXPECT_EQ(type->fields()[0].name(), "new_id"); + ICEBERG_UNWRAP_OR_FAIL(auto values, + ProjectPartitionValues(*type, *metadata_->partition_specs[0], + PartitionValues(Literal::Int(7)))); + ASSERT_EQ(values.num_fields(), 2); + EXPECT_EQ(values.values()[0], Literal::Long(7)); + EXPECT_TRUE(values.values()[1].IsNull()); +} + +} // namespace +} // namespace iceberg::internal