feat: add avro reader to registry

cmake_modules/IcebergThirdpartyToolchain.cmake

@@ -299,11 +299,26 @@ function(resolve_zlib_dependency)
 
 endfunction()
 
+# ----------------------------------------------------------------------
+# Zstd
+
+function(resolve_zstd_dependency)
+  find_package(zstd CONFIG)
+  if(zstd_FOUND)
+    list(APPEND ICEBERG_SYSTEM_DEPENDENCIES zstd)
+    message(STATUS "Found zstd, version: ${zstd_VERSION}")
+    set(ICEBERG_SYSTEM_DEPENDENCIES
+        ${ICEBERG_SYSTEM_DEPENDENCIES}
+        PARENT_SCOPE)
+  endif()
+endfunction()
+
 resolve_zlib_dependency()
 resolve_nanoarrow_dependency()
 resolve_nlohmann_json_dependency()
 
 if(ICEBERG_BUILD_BUNDLE)
   resolve_arrow_dependency()
   resolve_avro_dependency()
+  resolve_zstd_dependency()
 endif()

example/demo_example.cc

@@ -19,13 +19,16 @@
 
 #include <iostream>
 
-#include "iceberg/arrow/demo_arrow.h"
-#include "iceberg/avro/demo_avro.h"
-#include "iceberg/demo.h"
+#include "iceberg/avro/avro_reader.h"
+#include "iceberg/file_reader.h"
 
 int main() {
-  std::cout << iceberg::Demo().print() << std::endl;
-  std::cout << iceberg::arrow::DemoArrow().print() << std::endl;
-  std::cout << iceberg::avro::DemoAvro().print() << std::endl;
+  iceberg::avro::AvroReader::Register();
+  auto open_result = iceberg::ReaderFactoryRegistry::Open(
+      iceberg::FileFormatType::kAvro, {.path = "non-existing-file.avro"});
+  if (!open_result.has_value()) {
+    std::cerr << "Failed to open avro file" << std::endl;
+    return 1;
+  }
   return 0;
 }

src/iceberg/CMakeLists.txt

@@ -20,7 +20,6 @@ set(ICEBERG_INCLUDES "$<BUILD_INTERFACE:${PROJECT_BINARY_DIR}/src>"
 set(ICEBERG_SOURCES
     arrow_c_data_internal.cc
     catalog/in_memory_catalog.cc
-    demo.cc
     expression/expression.cc
     expression/literal.cc
     file_reader.cc
@@ -94,9 +93,7 @@ install(FILES ${CMAKE_CURRENT_BINARY_DIR}/iceberg_export.h
 
 if(ICEBERG_BUILD_BUNDLE)
   set(ICEBERG_BUNDLE_SOURCES
-      arrow/demo_arrow.cc
       arrow/arrow_fs_file_io.cc
-      avro/demo_avro.cc
       avro/avro_data_util.cc
       avro/avro_reader.cc
       avro/avro_schema_util.cc

src/iceberg/avro/avro_data_util.cc

@@ -419,14 +419,13 @@ Status AppendFieldToBuilder(const ::avro::NodePtr& avro_node,
                             const SchemaField& projected_field,
                             ::arrow::ArrayBuilder* array_builder) {
   if (avro_node->type() == ::avro::AVRO_UNION) {
-    const auto& union_datum = avro_datum.value<::avro::GenericUnion>();
-    size_t branch = union_datum.currentBranch();
+    size_t branch = avro_datum.unionBranch();
     if (avro_node->leafAt(branch)->type() == ::avro::AVRO_NULL) {
       ICEBERG_ARROW_RETURN_NOT_OK(array_builder->AppendNull());
       return {};
     } else {
-      return AppendFieldToBuilder(avro_node->leafAt(branch), union_datum.datum(),
-                                  projection, projected_field, array_builder);
+      return AppendFieldToBuilder(avro_node->leafAt(branch), avro_datum, projection,
+                                  projected_field, array_builder);
     }
   }
 

src/iceberg/avro/avro_reader.cc

@@ -76,9 +76,14 @@ struct ReadContext {
 // 1. prune the reader schema based on the projection
 // 2. read key-value metadata from the avro file
 // 3. collect basic reader metrics
-class AvroBatchReader::Impl {
+class AvroReader::Impl {
  public:
   Status Open(const ReaderOptions& options) {
+    // TODO(gangwu): perhaps adding a ReaderOptions::Validate() method
+    if (options.projection == nullptr) {
+      return InvalidArgument("Projected schema is required by Avro reader");
+    }
+
     batch_size_ = options.batch_size;
     read_schema_ = options.projection;
 
@@ -106,11 +111,10 @@ class AvroBatchReader::Impl {
 
     // Project the read schema on top of the file schema.
     // TODO(gangwu): support pruning source fields
-    ICEBERG_ASSIGN_OR_RAISE(projection_, Project(*options.projection, file_schema.root(),
+    ICEBERG_ASSIGN_OR_RAISE(projection_, Project(*read_schema_, file_schema.root(),
                                                  /*prune_source=*/false));
-    base_reader->init(file_schema);
     reader_ = std::make_unique<::avro::DataFileReader<::avro::GenericDatum>>(
-        std::move(base_reader));
+        std::move(base_reader), file_schema);
 
     if (options.split) {
       reader_->sync(options.split->offset);
@@ -119,7 +123,7 @@ class AvroBatchReader::Impl {
     return {};
   }
 
-  Result<Data> Next() {
+  Result<std::optional<ArrowArray>> Next() {
     if (!context_) {
       ICEBERG_RETURN_UNEXPECTED(InitReadContext());
     }
@@ -148,6 +152,19 @@ class AvroBatchReader::Impl {
     return {};
   }
 
+  Result<ArrowSchema> Schema() {
+    if (!context_) {
+      ICEBERG_RETURN_UNEXPECTED(InitReadContext());
+    }
+    ArrowSchema arrow_schema;
+    auto export_result = ::arrow::ExportSchema(*context_->arrow_schema_, &arrow_schema);
+    if (!export_result.ok()) {
+      return InvalidSchema("Failed to export the arrow schema: {}",
+                           export_result.message());
+    }
+    return arrow_schema;
+  }
+
  private:
   Status InitReadContext() {
     context_ = std::make_unique<ReadContext>();
@@ -174,9 +191,9 @@ class AvroBatchReader::Impl {
     return {};
   }
 
-  Result<Data> ConvertBuilderToArrowArray() {
+  Result<std::optional<ArrowArray>> ConvertBuilderToArrowArray() {
     if (context_->builder_->length() == 0) {
-      return {};
+      return std::nullopt;
     }
 
     auto builder_result = context_->builder_->Finish();
@@ -201,7 +218,7 @@ class AvroBatchReader::Impl {
   // The end of the split to read and used to terminate the reading.
   std::optional<int64_t> split_end_;
   // The schema to read.
-  std::shared_ptr<Schema> read_schema_;
+  std::shared_ptr<::iceberg::Schema> read_schema_;
   // The projection result to apply to the read schema.
   SchemaProjection projection_;
   // The avro reader to read the data into a datum.
@@ -210,13 +227,21 @@ class AvroBatchReader::Impl {
   std::unique_ptr<ReadContext> context_;
 };
 
-Result<Reader::Data> AvroBatchReader::Next() { return impl_->Next(); }
+Result<std::optional<ArrowArray>> AvroReader::Next() { return impl_->Next(); }
 
-Status AvroBatchReader::Open(const ReaderOptions& options) {
+Result<ArrowSchema> AvroReader::Schema() { return impl_->Schema(); }
+
+Status AvroReader::Open(const ReaderOptions& options) {
   impl_ = std::make_unique<Impl>();
   return impl_->Open(options);
 }
 
-Status AvroBatchReader::Close() { return impl_->Close(); }
+Status AvroReader::Close() { return impl_->Close(); }
+
+void AvroReader::Register() {
+  static ReaderFactoryRegistry avro_reader_register(
+      FileFormatType::kAvro,
+      []() -> Result<std::unique_ptr<Reader>> { return std::make_unique<AvroReader>(); });
+}
 
 }  // namespace iceberg::avro

src/iceberg/avro/avro_reader.h

@@ -25,19 +25,22 @@
 namespace iceberg::avro {
 
 /// \brief A reader that reads ArrowArray from Avro files.
-class ICEBERG_BUNDLE_EXPORT AvroBatchReader : public Reader {
+class ICEBERG_BUNDLE_EXPORT AvroReader : public Reader {
  public:
-  AvroBatchReader() = default;
+  AvroReader() = default;
 
-  ~AvroBatchReader() override = default;
+  ~AvroReader() override = default;
 
   Status Open(const ReaderOptions& options) final;
 
   Status Close() final;
 
-  Result<Data> Next() final;
+  Result<std::optional<ArrowArray>> Next() final;
 
-  DataLayout data_layout() const final { return DataLayout::kArrowArray; }
+  Result<ArrowSchema> Schema() final;
+
+  /// \brief Register this Avro reader implementation.
+  static void Register();
 
  private:
   class Impl;