a basic implement for export scan tasks as Arrow C ABI streams

Author: HeartLinked

src/iceberg/table_scan.cc

@@ -19,9 +19,13 @@
 
 #include "iceberg/table_scan.h"
 
-#include <algorithm>
-#include <ranges>
+#include <cstring>
+#include <vector>
 
+#include <iceberg/result.h>
+
+#include "iceberg/arrow_c_data.h"
+#include "iceberg/file_reader.h"
 #include "iceberg/manifest_entry.h"
 #include "iceberg/manifest_list.h"
 #include "iceberg/manifest_reader.h"
@@ -33,6 +37,112 @@
 
 namespace iceberg {
 
+namespace {
+/// \brief Private data structure to hold the Reader and error state
+struct ReaderStreamPrivateData {
+  std::unique_ptr<Reader> reader;
+  std::string last_error;
+
+  explicit ReaderStreamPrivateData(std::unique_ptr<Reader> reader_ptr)
+      : reader(std::move(reader_ptr)) {}
+
+  ~ReaderStreamPrivateData() {
+    if (reader) {
+      reader->Close();
+    }
+  }
+};
+
+/// \brief Callback to get the stream schema
+static int GetSchema(struct ArrowArrayStream* stream, struct ArrowSchema* out) {
+  if (!stream || !stream->private_data) {
+    return EINVAL;
+  }
+  auto* private_data = static_cast<ReaderStreamPrivateData*>(stream->private_data);
+  // Get schema from reader
+  auto schema_result = private_data->reader->Schema();
+  if (!schema_result.has_value()) {
+    private_data->last_error = schema_result.error().message;
+    std::memset(out, 0, sizeof(ArrowSchema));
+    return EIO;
+  }
+
+  *out = std::move(schema_result.value());
+  return 0;
+}
+
+/// \brief Callback to get the next array from the stream
+static int GetNext(struct ArrowArrayStream* stream, struct ArrowArray* out) {
+  if (!stream || !stream->private_data) {
+    return EINVAL;
+  }
+
+  auto* private_data = static_cast<ReaderStreamPrivateData*>(stream->private_data);
+
+  auto next_result = private_data->reader->Next();
+  if (!next_result.has_value()) {
+    private_data->last_error = next_result.error().message;
+    std::memset(out, 0, sizeof(ArrowArray));
+    return EIO;
+  }
+
+  auto& optional_array = next_result.value();
+  if (optional_array.has_value()) {
+    *out = std::move(optional_array.value());
+  } else {
+    // End of stream - set release to nullptr to signal end
+    std::memset(out, 0, sizeof(ArrowArray));
+    out->release = nullptr;
+  }
+
+  return 0;
+}
+
+/// \brief Callback to get the last error message
+static const char* GetLastError(struct ArrowArrayStream* stream) {
+  if (!stream || !stream->private_data) {
+    return nullptr;
+  }
+
+  auto* private_data = static_cast<ReaderStreamPrivateData*>(stream->private_data);
+  return private_data->last_error.empty() ? nullptr : private_data->last_error.c_str();
+}
+
+/// \brief Callback to release the stream resources
+static void Release(struct ArrowArrayStream* stream) {
+  if (!stream || !stream->private_data) {
+    return;
+  }
+
+  delete static_cast<ReaderStreamPrivateData*>(stream->private_data);
+  stream->private_data = nullptr;
+  stream->release = nullptr;
+}
+
+Result<ArrowArrayStream> MakeArrowArrayStream(std::unique_ptr<Reader> reader) {
+  if (!reader) {
+    return InvalidArgument("Reader cannot be null");
+  }
+
+  auto schema_check = reader->Schema();
+  if (!schema_check.has_value()) {
+    return InvalidSchema("Failed to get schema from reader: {}",
+                         schema_check.error().message);
+  }
+
+  auto private_data = std::make_unique<ReaderStreamPrivateData>(std::move(reader));
+
+  ArrowArrayStream stream{.get_schema = GetSchema,
+                          .get_next = GetNext,
+                          .get_last_error = GetLastError,
+                          .release = Release,
+                          .private_data = private_data.release()};
+
+  return stream;
+}
+
+}  // namespace
+
 // implement FileScanTask
 FileScanTask::FileScanTask(std::shared_ptr<DataFile> data_file)
     : data_file_(std::move(data_file)) {}
@@ -45,6 +155,21 @@ int32_t FileScanTask::files_count() const { return 1; }
 
 int64_t FileScanTask::estimated_row_count() const { return data_file_->record_count; }
 
+Result<ArrowArrayStream> FileScanTask::ToArrow(
+    const std::shared_ptr<Schema>& projected_schema,
+    const std::shared_ptr<Expression>& filter, const std::shared_ptr<FileIO>& io) const {
+  const ReaderOptions options{.path = data_file_->file_path,
+                              .length = data_file_->file_size_in_bytes,
+                              .io = io,
+                              .projection = projected_schema,
+                              .filter = filter};
+
+  ICEBERG_ASSIGN_OR_RAISE(auto reader,
+                          ReaderFactoryRegistry::Open(data_file_->file_format, options));
+
+  return MakeArrowArrayStream(std::move(reader));
+}
+
 TableScanBuilder::TableScanBuilder(std::shared_ptr<TableMetadata> table_metadata,
                                    std::shared_ptr<FileIO> file_io)
     : file_io_(std::move(file_io)) {
@@ -178,4 +303,28 @@ Result<std::vector<std::shared_ptr<FileScanTask>>> DataTableScan::PlanFiles() co
   return tasks;
 }
 
+Result<std::vector<ArrowArrayStream>> DataTableScan::ToArrow() const {
+  Result<std::vector<std::shared_ptr<FileScanTask>>> tasks_result = PlanFiles();
+  if (!tasks_result.has_value()) {
+    return InvalidArgument("Failed to plan files: {}", tasks_result.error().message);
+  }
+  auto tasks = tasks_result.value();
+  if (tasks.empty()) {
+    // TODO(Li Feiyang): return a empty arrow stream
+    return NotImplemented("No files to scan");
+  }
+
+  std::vector<ArrowArrayStream> arrow_streams;
+  for (const auto& task : tasks_result.value()) {
+    Result<ArrowArrayStream> arrow_stream_result =
+        task->ToArrow(context_.projected_schema, context_.filter, file_io_);
+    if (!arrow_stream_result.has_value()) {
+      return InvalidArgument("Failed to get arrow stream: {}",
+                             arrow_stream_result.error().message);
+    }
+    arrow_streams.push_back(arrow_stream_result.value());
+  }
+  return std::move(arrow_streams);
+}
+
 }  // namespace iceberg

src/iceberg/table_scan.h

@@ -22,6 +22,7 @@
 #include <string>
 #include <vector>
 
+#include "iceberg/arrow_c_data.h"
 #include "iceberg/manifest_entry.h"
 #include "iceberg/type_fwd.h"
 
@@ -54,6 +55,11 @@ class ICEBERG_EXPORT FileScanTask : public ScanTask {
   int32_t files_count() const override;
   int64_t estimated_row_count() const override;
 
+  Result<ArrowArrayStream> ToArrow(const std::shared_ptr<Schema>& projected_schema,
+
+                                   const std::shared_ptr<Expression>& filter,
+                                   const std::shared_ptr<FileIO>& io) const;
+
  private:
   /// \brief Data file metadata.
   std::shared_ptr<DataFile> data_file_;
@@ -183,6 +189,8 @@ class ICEBERG_EXPORT DataTableScan : public TableScan {
   /// \brief Plans the scan tasks by resolving manifests and data files.
   /// \return A Result containing scan tasks or an error.
   Result<std::vector<std::shared_ptr<FileScanTask>>> PlanFiles() const override;
+
+  Result<std::vector<ArrowArrayStream>> ToArrow() const;
 };
 
 }  // namespace iceberg