TODO

Author: Yin Ye

tensorflow_io/core/kernels/arrow/arrow_dataset_ops.cc

@@ -32,8 +32,10 @@ namespace data {
 
 namespace {
 
+// TODO(yye): implement this func and support reading parquet data in
+// streamining way.
 Status OpenParquetFile(arrow::fs::FileSystem *fs,
-                       arrow::io::RandomAccessFile *file, )
+                       arrow::io::RandomAccessFile *file) {}
 
 }  // namespace
 
@@ -616,55 +618,37 @@ class ArrowParquetDatasetOp : public ArrowOpKernelBase {
           : ArrowBaseIterator<Dataset>(params) {}
 
      private:
+      // TODO(yye): implementation of getting the first batch.
       Status SetupStreamsLocked(Env *env)
           TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) override {
         TF_RETURN_IF_ERROR(ReadFile(current_file_idx_));
 
+        // Open and read parquet file.
         while (record_batches_.empty() &&
                ++current_file_idx_ < dataset()->file_paths_.size()) {
           TF_RETURN_IF_ERROR(ReadFile(current_file_idx_));
         }
 
-        if (current_batch_idx_ < record_batches_.size()) {
-          current_batch_ = record_batches_[current_batch_idx_];
-        }
         return OkStatus();
       }
 
+      // TODO(yye): implementation of getting the next batch.
       Status NextStreamLocked(Env *env)
           TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) override {
-        ArrowBaseIterator<Dataset>::NextStreamLocked(env);
-        if (++current_batch_idx_ < record_batches_.size()) {
-          current_batch_ = record_batches_[current_batch_idx_];
-        } else if (++current_file_idx_ < dataset()->file_paths_.size()) {
-          current_batch_idx_ = 0;
-          while (record_batches_.empty() &&
-                 ++current_file_idx_ < dataset()->file_paths_.size()) {
-            TF_RETURN_IF_ERROR(ReadFile(current_file_idx_));
-          }
-
-          if (!record_batches_.empty()) {
-            current_batch_ = record_batches_[current_batch_idx_];
-          } else {
-            current_batch_ = nullptr;
-          }
-        }
         return OkStatus();
       }
 
       void ResetStreamsLocked() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) override {
         ArrowBaseIterator<Dataset>::ResetStreamsLocked();
         current_file_idx_ = 0;
-        current_batch_idx_ = 0;
         record_batches_.clear();
       }
 
       Status ReadFile(int file_index) TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
         Status res = OkStatus();
         do {
           std::shared_ptr<arrow::io::RandomAccessFile> file;
-          res = ArrowUtil::OpenParquetFile(&fs_, &file,
-                                           dataset()->file_paths_[file_index]);
+          res = OpenParquetFile(&fs_, file.get());
           if (!res.ok()) {
             break;
           }
@@ -682,7 +666,7 @@ class ArrowParquetDatasetOp : public ArrowOpKernelBase {
           std::unique_ptr<parquet::arrow::FileReader> reader;
           builder->properties(properties)->Build(&reader);
 
-          if (column_indices_.empty() || !dataset()->same_header_) {
+          if (column_indices_.empty()) {
             column_indices_.clear();
             std::shared_ptr<arrow::Schema> schema;
             reader->GetSchema(&schema);
@@ -741,9 +725,12 @@ class ArrowParquetDatasetOp : public ArrowOpKernelBase {
       }
 
       size_t current_file_idx_ TF_GUARDED_BY(mu_) = 0;
+      // TODO(yye): stop maintaining/holding all the record batches.
       std::vector<std::shared_ptr<arrow::RecordBatch>> record_batches_
           TF_GUARDED_BY(mu_);
       std::shared_ptr<arrow::fs::FileSystem> fs_ TF_GUARDED_BY(mu_) = nullptr;
+
+      // Maintains the index of the columns to read.
       std::vector<int> column_indices_ TF_GUARDED_BY(mu_);
     };