Revert "Avoid whole partition cache (#38)" (#40)

This reverts commit 8d7e250.

cpp/src/arrow/dataset/file_parquet.cc

@@ -51,7 +51,6 @@ using parquet::arrow::SchemaField;
 using parquet::arrow::SchemaManifest;
 using parquet::arrow::StatisticsAsScalars;
 
-
 /// \brief A ScanTask backed by a parquet file and a RowGroup within a parquet file.
 class ParquetScanTask : public ScanTask {
  public:

cpp/src/arrow/dataset/scanner.cc

@@ -22,7 +22,6 @@
 #include <memory>
 #include <mutex>
 #include <sstream>
-#include <iostream>
 
 #include "arrow/array/array_primitive.h"
 #include "arrow/compute/api_scalar.h"
@@ -268,7 +267,7 @@ class ARROW_DS_EXPORT SyncScanner : public Scanner {
       : Scanner(std::move(scan_options)), fragment_(std::move(fragment)) {}
 
   Result<TaggedRecordBatchIterator> ScanBatches() override;
-  Result<RecordBatchIterator> ScanBatchesWithWeakFilterProject() override;
+  Result<TaggedRecordBatchIterator> ScanBatchesWithWeakFilterProject() override;
   Result<ScanTaskIterator> Scan() override;
   Status Scan(std::function<Status(TaggedRecordBatch)> visitor) override;
   Result<std::shared_ptr<Table>> ToTable() override;
@@ -299,23 +298,28 @@ Result<TaggedRecordBatchIterator> SyncScanner::ScanBatches() {
   });
 }
 
-Result<RecordBatchIterator> SyncScanner::ScanBatchesWithWeakFilterProject() {
+Result<TaggedRecordBatchIterator> SyncScanner::ScanBatchesWithWeakFilterProject() {
   ARROW_ASSIGN_OR_RAISE(auto fragment_it, GetFragments())
   auto fn = [this](const std::shared_ptr<Fragment>& fragment) -> Result<ScanTaskIterator> {
     ARROW_ASSIGN_OR_RAISE(auto scan_task_it, fragment->Scan(scan_options_));
     return std::move(scan_task_it);
   };
+
   // Iterator<Iterator<ScanTask>>
   auto maybe_scantask_it = MakeMaybeMapIterator(fn, std::move(fragment_it));
   auto scan_task_it = MakeFlattenIterator(std::move(maybe_scantask_it));
-
-  auto scan_fn = [](const std::shared_ptr<ScanTask>& aTask)->Result<RecordBatchIterator> {
-    ARROW_ASSIGN_OR_RAISE(auto recordbatch_it, aTask->Execute());
-    return std::move(recordbatch_it);    
-  };
 
-  auto maybe_recordbatch_it = MakeMaybeMapIterator(scan_fn, std::move(scan_task_it));
-  return MakeFlattenIterator(std::move(maybe_recordbatch_it));
+  auto task_group = scan_options_->TaskGroup();
+  auto state = std::make_shared<ScanBatchesState>(std::move(scan_task_it), task_group);
+  for (int i = 0; i < scan_options_->fragment_readahead; i++) {
+    state->PushScanTask();
+  }
+  return MakeFunctionIterator([task_group, state]() -> Result<TaggedRecordBatch> {
+    ARROW_ASSIGN_OR_RAISE(auto batch, state->Pop());
+    if (!IsIterationEnd(batch)) return batch;
+    RETURN_NOT_OK(task_group->Finish());
+    return IterationEnd<TaggedRecordBatch>();
+  });
 }
 
 Result<FragmentIterator> SyncScanner::GetFragments() {
@@ -381,7 +385,7 @@ class ARROW_DS_EXPORT AsyncScanner : public Scanner,
 
   Status Scan(std::function<Status(TaggedRecordBatch)> visitor) override;
   Result<TaggedRecordBatchIterator> ScanBatches() override;
-  Result<RecordBatchIterator> ScanBatchesWithWeakFilterProject() override;
+  Result<TaggedRecordBatchIterator> ScanBatchesWithWeakFilterProject() override;
   Result<EnumeratedRecordBatchIterator> ScanBatchesUnordered() override;
   Result<std::shared_ptr<Table>> ToTable() override;
 
@@ -504,7 +508,7 @@ Result<TaggedRecordBatchIterator> AsyncScanner::ScanBatches() {
   return MakeGeneratorIterator(std::move(batches_gen));
 }
 
-Result<RecordBatchIterator> AsyncScanner::ScanBatchesWithWeakFilterProject() {
+Result<TaggedRecordBatchIterator> AsyncScanner::ScanBatchesWithWeakFilterProject() {
   return Status::NotImplemented("Scanning with weak filter project not implemented in async scanner");
 }
 

cpp/src/arrow/dataset/scanner.h

@@ -277,7 +277,7 @@ class ARROW_DS_EXPORT Scanner {
   ///
   /// Filter and Project expressions in ScanOption will be not executed individually and
   /// will be only recognized by the underlying file format.
-  virtual Result<RecordBatchIterator> ScanBatchesWithWeakFilterProject() = 0;
+  virtual Result<TaggedRecordBatchIterator> ScanBatchesWithWeakFilterProject() = 0;
   /// \brief Scan the dataset into a stream of record batches.  Unlike ScanBatches this
   /// method may allow record batches to be returned out of order.  This allows for more
   /// efficient scanning: some fragments may be accessed more quickly than others (e.g.

cpp/src/jni/dataset/jni_wrapper.cc

@@ -125,7 +125,7 @@ class ReserveFromJava : public arrow::jniutil::ReservationListener {
 class DisposableScannerAdaptor {
  public:
   DisposableScannerAdaptor(std::shared_ptr<arrow::dataset::Scanner> scanner,
-                           arrow::RecordBatchIterator batch_itr)
+                           arrow::dataset::TaggedRecordBatchIterator batch_itr)
       : scanner_(std::move(scanner)), batch_itr_(std::move(batch_itr)) {}
 
   static arrow::Result<std::shared_ptr<DisposableScannerAdaptor>> Create(
@@ -136,18 +136,18 @@ class DisposableScannerAdaptor {
 
   arrow::Result<std::shared_ptr<arrow::RecordBatch>> Next() {
     ARROW_ASSIGN_OR_RAISE(std::shared_ptr<arrow::RecordBatch> batch, NextBatch());
-    return std::move(batch);
+    return batch;
   }
 
   const std::shared_ptr<arrow::dataset::Scanner>& GetScanner() const { return scanner_; }
 
  private:
   std::shared_ptr<arrow::dataset::Scanner> scanner_;
-  arrow::RecordBatchIterator batch_itr_;
+  arrow::dataset::TaggedRecordBatchIterator batch_itr_;
 
   arrow::Result<std::shared_ptr<arrow::RecordBatch>> NextBatch() {
     ARROW_ASSIGN_OR_RAISE(auto batch, batch_itr_.Next())
-    return batch;
+    return batch.record_batch;
   }
 };