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/test/CMakeLists.txt b/src/iceberg/test/CMakeLists.txt index 1874fc25e..ec4721635 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_rows_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_rows_test.cc b/src/iceberg/test/inspect/metadata_rows_test.cc new file mode 100644 index 000000000..be69880a0 --- /dev/null +++ b/src/iceberg/test/inspect/metadata_rows_test.cc @@ -0,0 +1,137 @@ +/* + * 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_util_internal.h" +#include "iceberg/inspect/metadata_table.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); +} + +} // namespace +} // namespace iceberg::internal