Simplify RecordBatchReader to use Project.next()

cpp/src/lance/arrow/file_lance.cc

@@ -63,18 +63,8 @@ ::arrow::Result<::arrow::RecordBatchGenerator> LanceFileFormat::ScanBatchesAsync
   auto reader = std::make_shared<lance::io::FileReader>(infile);
   ARROW_RETURN_NOT_OK(reader->Open());
 
-  std::optional<int64_t> limit = std::nullopt;
-  int64_t offset = 0;
-  if (options->fragment_scan_options &&
-      options->fragment_scan_options->type_name() == kLanceFormatTypeName) {
-    auto lance_fragment_scan_options =
-        std::dynamic_pointer_cast<LanceFragmentScanOptions>(options->fragment_scan_options);
-    limit = lance_fragment_scan_options->limit;
-    offset = lance_fragment_scan_options->offset;
-  }
-
-  auto batch_reader = lance::io::RecordBatchReader(
-      reader, options, ::arrow::internal::GetCpuThreadPool(), limit, offset);
+  auto batch_reader =
+      lance::io::RecordBatchReader(reader, options, ::arrow::internal::GetCpuThreadPool());
   ARROW_RETURN_NOT_OK(batch_reader.Open());
   auto generator = ::arrow::RecordBatchGenerator(std::move(batch_reader));
   return generator;

cpp/src/lance/io/exec/filter.h

@@ -14,7 +14,7 @@
 
 #pragma once
 
-#include <arrow/compute/exec/expression.h>
+#include <arrow/compute/api.h>
 #include <arrow/result.h>
 
 #include <memory>

cpp/src/lance/io/exec/project.cc

@@ -18,6 +18,7 @@
 #include <arrow/result.h>
 #include <fmt/format.h>
 
+#include "lance/arrow/file_lance_ext.h"
 #include "lance/arrow/utils.h"
 #include "lance/io/exec/filter.h"
 #include "lance/io/exec/limit.h"
@@ -31,16 +32,23 @@ Project::Project(std::unique_ptr<ExecNode> child) : child_(std::move(child)) {}
 
 ::arrow::Result<std::unique_ptr<Project>> Project::Make(
     std::shared_ptr<FileReader> reader,
-    std::shared_ptr<::arrow::dataset::ScanOptions> scan_options,
-    std::optional<int32_t> limit,
-    int32_t offset) {
+    std::shared_ptr<::arrow::dataset::ScanOptions> scan_options) {
   auto& schema = reader->schema();
   auto projected_arrow_schema = scan_options->projected_schema;
   if (projected_arrow_schema->num_fields() == 0) {
     projected_arrow_schema = scan_options->dataset_schema;
   }
   ARROW_ASSIGN_OR_RAISE(auto projected_schema, schema.Project(*projected_arrow_schema));
 
+  std::optional<int64_t> limit;
+  int64_t offset;
+  if (scan_options->fragment_scan_options &&
+      scan_options->fragment_scan_options->type_name() == "lance") {
+    auto fso = std::dynamic_pointer_cast<lance::arrow::LanceFragmentScanOptions>(
+        scan_options->fragment_scan_options);
+    limit = fso->limit;
+    offset = fso->offset;
+  }
   std::unique_ptr<ExecNode> child;
   if (Filter::HasFilter(scan_options->filter)) {
     /// Build a chain of:

cpp/src/lance/io/exec/project.h

@@ -45,9 +45,7 @@ class Project : ExecNode {
   ///
   static ::arrow::Result<std::unique_ptr<Project>> Make(
       std::shared_ptr<FileReader> reader,
-      std::shared_ptr<::arrow::dataset::ScanOptions> scan_options,
-      std::optional<int32_t> limit = std::nullopt,
-      int32_t offset = 0);
+      std::shared_ptr<::arrow::dataset::ScanOptions> scan_options);
 
   ::arrow::Result<ScanBatch> Next() override;
 

cpp/src/lance/io/record_batch_reader.cc

@@ -36,38 +36,25 @@ namespace lance::io {
 
 RecordBatchReader::RecordBatchReader(std::shared_ptr<FileReader> reader,
                                      std::shared_ptr<::arrow::dataset::ScanOptions> options,
-                                     ::arrow::internal::ThreadPool* thread_pool,
-                                     std::optional<int64_t> limit,
-                                     int64_t offset) noexcept
-    : reader_(reader),
-      options_(options),
-      limit_(limit),
-      offset_(offset),
-      thread_pool_(thread_pool) {
+                                     ::arrow::internal::ThreadPool* thread_pool) noexcept
+    : reader_(reader), options_(options), thread_pool_(thread_pool) {
   assert(thread_pool_);
 }
 
 RecordBatchReader::RecordBatchReader(const RecordBatchReader& other) noexcept
     : reader_(other.reader_),
       options_(other.options_),
-      limit_(other.limit_),
-      offset_(other.offset_),
       project_(other.project_),
-      thread_pool_(other.thread_pool_),
-      current_batch_(int(other.current_batch_)) {}
+      thread_pool_(other.thread_pool_) {}
 
 RecordBatchReader::RecordBatchReader(RecordBatchReader&& other) noexcept
     : reader_(std::move(other.reader_)),
       options_(std::move(other.options_)),
-      limit_(other.limit_),
-      offset_(other.offset_),
       project_(std::move(other.project_)),
-      thread_pool_(std::move(other.thread_pool_)),
-      current_batch_(int(other.current_batch_)),
-      readahead_queue_(std::move(other.readahead_queue_)) {}
+      thread_pool_(std::move(other.thread_pool_)) {}
 
 ::arrow::Status RecordBatchReader::Open() {
-  ARROW_ASSIGN_OR_RAISE(project_, exec::Project::Make(reader_, options_, limit_, offset_));
+  ARROW_ASSIGN_OR_RAISE(project_, exec::Project::Make(reader_, options_));
   return ::arrow::Status::OK();
 }
 
@@ -81,34 +68,19 @@ ::arrow::Status RecordBatchReader::ReadNext(std::shared_ptr<::arrow::RecordBatch
   return ::arrow::Status::OK();
 }
 
-::arrow::Result<std::shared_ptr<::arrow::RecordBatch>> RecordBatchReader::ReadBatch(
-    [[maybe_unused]] int32_t batch_id) const {
+::arrow::Result<std::shared_ptr<::arrow::RecordBatch>> RecordBatchReader::ReadBatch() const {
   ARROW_ASSIGN_OR_RAISE(auto batch, project_->Next());
   return batch.batch;
 }
 
 ::arrow::Future<std::shared_ptr<::arrow::RecordBatch>> RecordBatchReader::operator()() {
-  int total_batches = reader_->metadata().num_batches();
-  while (static_cast<int32_t>(readahead_queue_.size()) < options_->batch_readahead &&
-         current_batch_ < total_batches) {
-    int32_t batch_id = current_batch_++;
-    auto result = thread_pool_->Submit(
-        [&](int32_t batch_id) {
-          return ::arrow::Future<std::shared_ptr<::arrow::RecordBatch>>::MakeFinished(
-              this->ReadBatch(batch_id));
-        },
-        batch_id);
-    if (!result.ok()) {
-      return result.status();
-    }
-    readahead_queue_.emplace(std::move(result.ValueOrDie()));
+  auto result = thread_pool_->Submit([&]() {
+    return ::arrow::Future<std::shared_ptr<::arrow::RecordBatch>>::MakeFinished(this->ReadBatch());
+  });
+  if (!result.ok()) {
+    return result.status();
   }
-  if (readahead_queue_.empty()) {
-    return std::shared_ptr<::arrow::RecordBatch>(nullptr);
-  }
-  auto first = readahead_queue_.front();
-  readahead_queue_.pop();
-  return first;
+  return ::arrow::Future<std::shared_ptr<::arrow::RecordBatch>>(result.ValueOrDie());
 }
 
 }  // namespace lance::io

cpp/src/lance/io/record_batch_reader.h

@@ -46,9 +46,7 @@ class RecordBatchReader : ::arrow::RecordBatchReader {
   /// Constructor.
   RecordBatchReader(std::shared_ptr<FileReader> reader,
                     std::shared_ptr<::arrow::dataset::ScanOptions> options,
-                    ::arrow::internal::ThreadPool* thread_pool_,
-                    std::optional<int64_t> limit = std::nullopt,
-                    int64_t offset = 0) noexcept;
+                    ::arrow::internal::ThreadPool* thread_pool_) noexcept;
 
   /// Copy constructor.
   RecordBatchReader(const RecordBatchReader& other) noexcept;
@@ -71,18 +69,13 @@ class RecordBatchReader : ::arrow::RecordBatchReader {
  private:
   RecordBatchReader() = delete;
 
-  ::arrow::Result<std::shared_ptr<::arrow::RecordBatch>> ReadBatch(int32_t batch_id) const;
+  ::arrow::Result<std::shared_ptr<::arrow::RecordBatch>> ReadBatch() const;
 
   std::shared_ptr<FileReader> reader_;
   std::shared_ptr<::arrow::dataset::ScanOptions> options_;
-  std::optional<int64_t> limit_ = std::nullopt;
-  int64_t offset_ = 0;
   /// Projection over the dataset.
   std::shared_ptr<exec::Project> project_;
-
   ::arrow::internal::ThreadPool* thread_pool_;
-  std::atomic_int32_t current_batch_ = 0;
-  std::queue<::arrow::Future<std::shared_ptr<::arrow::RecordBatch>>> readahead_queue_;
 };
 
 }  // namespace lance::io