// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. #pragma once #include <memory> #include "arrow/python/common.h" #include "arrow/python/visibility.h" #include "arrow/record_batch.h" #include "arrow/result.h" #include "arrow/util/macros.h" namespace arrow { namespace py { class ARROW_PYTHON_EXPORT PyRecordBatchReader : public RecordBatchReader { public: std::shared_ptr<Schema> schema() const override; Status ReadNext(std::shared_ptr<RecordBatch>* batch) override; // For use from Cython // Assumes that `iterable` is borrowed static Result<std::shared_ptr<RecordBatchReader>> Make(std::shared_ptr<Schema>, PyObject* iterable); protected: PyRecordBatchReader(); Status Init(std::shared_ptr<Schema>, PyObject* iterable); std::shared_ptr<Schema> schema_; OwnedRefNoGIL iterator_; }; class ARROW_PYTHON_EXPORT CastingRecordBatchReader : public RecordBatchReader { public: std::shared_ptr<Schema> schema() const override; Status ReadNext(std::shared_ptr<RecordBatch>* batch) override; static Result<std::shared_ptr<RecordBatchReader>> Make( std::shared_ptr<RecordBatchReader> parent, std::shared_ptr<Schema> schema); Status Close() override; protected: CastingRecordBatchReader(); Status Init(std::shared_ptr<RecordBatchReader> parent, std::shared_ptr<Schema> schema); std::shared_ptr<RecordBatchReader> parent_; std::shared_ptr<Schema> schema_; }; } // namespace py } // namespace arrow
Memory