| 140 | std::shared_ptr<arrow::Schema> schema() const override { return schema_; } |
| 141 | |
| 142 | arrow::Status ReadNext(std::shared_ptr<arrow::RecordBatch>* batch_out) override { |
| 143 | if (done_) { |
| 144 | // Close() has been called |
| 145 | batch_out = nullptr; |
| 146 | return arrow::Status::OK(); |
| 147 | } |
| 148 | |
| 149 | ARROW_RETURN_NOT_OK(reader_->ReadNext(batch_out)); |
| 150 | if (batch_out->get()) { |
| 151 | num_rows_ -= batch_out->get()->num_rows(); |
| 152 | if (num_rows_ < 0) { |
| 153 | auto smaller_batch = |
| 154 | batch_out->get()->Slice(0, batch_out->get()->num_rows() + num_rows_); |
| 155 | *batch_out = smaller_batch; |
| 156 | } |
| 157 | |
| 158 | if (num_rows_ <= 0) { |
| 159 | // We've run out of num_rows before batches |
| 160 | ARROW_RETURN_NOT_OK(Close()); |
| 161 | } |
| 162 | } else { |
| 163 | // We've run out of batches before num_rows |
| 164 | ARROW_RETURN_NOT_OK(Close()); |
| 165 | } |
| 166 | |
| 167 | return arrow::Status::OK(); |
| 168 | } |
| 169 | |
| 170 | arrow::Status Close() override { |
| 171 | if (done_) { |