refactor MakeReader

Author: xiao.dong

src/iceberg/manifest_reader.cc

@@ -19,18 +19,40 @@
 
 #include "iceberg/manifest_reader.h"
 
+#include "iceberg/manifest_entry.h"
+#include "iceberg/manifest_list.h"
 #include "iceberg/manifest_reader_internal.h"
+#include "iceberg/schema.h"
+#include "iceberg/util/macros.h"
 
 namespace iceberg {
 
-std::shared_ptr<ManifestReader> ManifestReader::NewReader(
-    std::unique_ptr<Reader> reader) {
-  return std::make_shared<ManifestReaderImpl>(std::move(reader));
+Result<std::shared_ptr<ManifestReader>> ManifestReader::MakeReader(
+    const std::string& manifest_location, std::shared_ptr<FileIO> file_io,
+    std::shared_ptr<Schema> partition_schema) {
+  auto manifest_entry_schema = ManifestEntry::TypeFromPartitionType(partition_schema);
+  auto fields_span = manifest_entry_schema->fields();
+  std::vector<SchemaField> fields(fields_span.begin(), fields_span.end());
+  auto schema = std::make_shared<Schema>(fields);
+  ICEBERG_ASSIGN_OR_RAISE(
+      auto reader,
+      ReaderFactoryRegistry::Open(
+          FileFormatType::kAvro,
+          {.path = manifest_location, .io = std::move(file_io), .projection = schema}));
+  return std::make_shared<ManifestReaderImpl>(std::move(reader), std::move(schema));
 }
 
-std::shared_ptr<ManifestListReader> ManifestListReader::NewReader(
-    std::unique_ptr<Reader> reader) {
-  return std::make_shared<ManifestListReaderImpl>(std::move(reader));
+Result<std::shared_ptr<ManifestListReader>> ManifestListReader::MakeReader(
+    const std::string& manifest_list_location, std::shared_ptr<FileIO> file_io) {
+  std::vector<SchemaField> fields(ManifestFile::Type().fields().begin(),
+                                  ManifestFile::Type().fields().end());
+  auto schema = std::make_shared<Schema>(fields);
+  ICEBERG_ASSIGN_OR_RAISE(
+      auto reader,
+      ReaderFactoryRegistry::Open(FileFormatType::kAvro, {.path = manifest_list_location,
+                                                          .io = std::move(file_io),
+                                                          .projection = schema}));
+  return std::make_shared<ManifestListReaderImpl>(std::move(reader), std::move(schema));
 }
 
 }  // namespace iceberg

src/iceberg/manifest_reader.h

@@ -36,15 +36,18 @@ class ICEBERG_EXPORT ManifestReader {
  public:
   virtual Result<std::vector<std::unique_ptr<ManifestEntry>>> Entries() const = 0;
 
-  static std::shared_ptr<ManifestReader> NewReader(std::unique_ptr<Reader> reader);
+  static Result<std::shared_ptr<ManifestReader>> MakeReader(
+      const std::string& manifest_location, std::shared_ptr<FileIO> file_io,
+      std::shared_ptr<Schema> partition_schema);
 };
 
 /// \brief Read manifest files from a manifest list file.
 class ICEBERG_EXPORT ManifestListReader {
  public:
   virtual Result<std::vector<std::unique_ptr<ManifestFile>>> Files() const = 0;
 
-  static std::shared_ptr<ManifestListReader> NewReader(std::unique_ptr<Reader> reader);
+  static Result<std::shared_ptr<ManifestListReader>> MakeReader(
+      const std::string& manifest_list_location, std::shared_ptr<FileIO> file_io);
 };
 
 }  // namespace iceberg

src/iceberg/manifest_reader_internal.cc

@@ -29,12 +29,13 @@
 #include "iceberg/schema.h"
 #include "iceberg/schema_internal.h"
 #include "iceberg/type.h"
+#include "iceberg/util/macros.h"
 
 namespace iceberg {
 
 #define ARROW_RETURN_IF_NOT_OK(status, error)                      \
-  if (status != NANOARROW_OK) {                                    \
-    return InvalidArrowData("NanoArrow error: {}", error.message); \
+  if (status != NANOARROW_OK) [[unlikely]] {                       \
+    return InvalidArrowData("Nanoarrow error: {}", error.message); \
   }
 
 Result<std::vector<std::unique_ptr<ManifestFile>>> ParseManifestListEntry(
@@ -206,16 +207,8 @@ Result<std::vector<std::unique_ptr<ManifestEntry>>> ManifestReaderImpl::Entries(
 
 Result<std::vector<std::unique_ptr<ManifestFile>>> ManifestListReaderImpl::Files() const {
   std::vector<std::unique_ptr<ManifestFile>> manifest_files;
-  auto arrow_schema = reader_->Schema();
-  if (!arrow_schema.has_value()) {
-    return InvalidArgument("Get schema failed in reader:{}",
-                           arrow_schema.error().message);
-  }
-  internal::ArrowSchemaGuard schema_guard(&arrow_schema.value());
-  auto schema = FromArrowSchema(arrow_schema.value(), std::nullopt);
-  if (!schema.has_value()) {
-    return InvalidArgument("Parse iceberg schema failed:{}", schema.error().message);
-  }
+  ICEBERG_ASSIGN_OR_RAISE(auto arrow_schema, reader_->Schema());
+  internal::ArrowSchemaGuard schema_guard(&arrow_schema);
   while (true) {
     auto result = reader_->Next();
     if (!result.has_value()) {
@@ -224,8 +217,8 @@ Result<std::vector<std::unique_ptr<ManifestFile>>> ManifestListReaderImpl::Files
     }
     if (result.value().has_value()) {
       internal::ArrowArrayGuard array_guard(&result.value().value());
-      auto parse_result = ParseManifestListEntry(
-          &arrow_schema.value(), &result.value().value(), *schema.value());
+      auto parse_result =
+          ParseManifestListEntry(&arrow_schema, &result.value().value(), *schema_);
       if (!parse_result.has_value()) {
         return InvalidArgument("Failed to parse manifest list entry:{}",
                                parse_result.error().message);

src/iceberg/manifest_reader_internal.h

@@ -29,24 +29,28 @@ namespace iceberg {
 /// \brief Read manifest entries from a manifest file.
 class ManifestReaderImpl : public ManifestReader {
  public:
-  explicit ManifestReaderImpl(std::unique_ptr<Reader> reader)
-      : reader_(std::move(reader)) {}
+  explicit ManifestReaderImpl(std::unique_ptr<Reader> reader,
+                              std::shared_ptr<Schema> schema)
+      : schema_(std::move(schema)), reader_(std::move(reader)) {}
 
   Result<std::vector<std::unique_ptr<ManifestEntry>>> Entries() const override;
 
  private:
+  std::shared_ptr<Schema> schema_;
   std::unique_ptr<Reader> reader_;
 };
 
 /// \brief Read manifest files from a manifest list file.
 class ManifestListReaderImpl : public ManifestListReader {
  public:
-  explicit ManifestListReaderImpl(std::unique_ptr<Reader> reader)
-      : reader_(std::move(reader)) {}
+  explicit ManifestListReaderImpl(std::unique_ptr<Reader> reader,
+                                  std::shared_ptr<Schema> schema)
+      : schema_(std::move(schema)), reader_(std::move(reader)) {}
 
   Result<std::vector<std::unique_ptr<ManifestFile>>> Files() const override;
 
  private:
+  std::shared_ptr<Schema> schema_;
   std::unique_ptr<Reader> reader_;
 };
 

test/manifest_list_reader_test.cc

@@ -47,16 +47,11 @@ class ManifestListReaderTest : public TempFileTestBase {
 };
 
 TEST_F(ManifestListReaderTest, BasicTest) {
-  std::vector<SchemaField> fields(ManifestFile::Type().fields().begin(),
-                                  ManifestFile::Type().fields().end());
-  auto schema = std::make_shared<Schema>(fields);
   std::string path = GetResourcePath(
       "snap-7412193043800610213-1-2bccd69e-d642-4816-bba0-261cd9bd0d93.avro");
-  auto reader_result = ReaderFactoryRegistry::Open(
-      FileFormatType::kAvro, {.path = path, .io = file_io_, .projection = schema});
-  ASSERT_THAT(reader_result, IsOk());
-  auto reader = std::move(reader_result.value());
-  auto manifest_reader = ManifestListReader::NewReader(std::move(reader));
+  auto manifest_reader_result = ManifestListReader::MakeReader(path, file_io_);
+  ASSERT_EQ(manifest_reader_result.has_value(), true);
+  auto manifest_reader = manifest_reader_result.value();
   auto read_result = manifest_reader->Files();
   ASSERT_EQ(read_result.has_value(), true);
   ASSERT_EQ(read_result.value().size(), 4);