arrow_array_stream_proxy.hpp

Go to the documentation of this file.

// Copyright 2024 Man Group Operations Limited
//
// Licensed 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 <ranges>
 
#include "sparrow/array.hpp"
#include "sparrow/array_api.hpp"
#include "sparrow/arrow_interface/arrow_array.hpp"
#include "sparrow/arrow_interface/arrow_array_stream.hpp"
#include "sparrow/arrow_interface/arrow_schema.hpp"
#include "sparrow/c_interface.hpp"
#include "sparrow/layout/layout_concept.hpp"
 
namespace sparrow
{
    class arrow_array_stream_proxy
    {
    public:

        SPARROW_API arrow_array_stream_proxy();

        SPARROW_API explicit arrow_array_stream_proxy(ArrowArrayStream&& stream);

        SPARROW_API explicit arrow_array_stream_proxy(ArrowArrayStream* stream);
 
        // explicit arrow_array_stream_proxy(ArrowSchema* schema_ptr);
 
        arrow_array_stream_proxy(const arrow_array_stream_proxy&) = delete;
        arrow_array_stream_proxy& operator=(const arrow_array_stream_proxy&) = delete;
 
        SPARROW_API
        arrow_array_stream_proxy(arrow_array_stream_proxy&& other) noexcept;
        SPARROW_API
        arrow_array_stream_proxy& operator=(arrow_array_stream_proxy&& other) noexcept;

        SPARROW_API ~arrow_array_stream_proxy();

        [[nodiscard]] SPARROW_API bool owns_stream() const;

        [[nodiscard]] SPARROW_API ArrowArrayStream* export_stream();

        template <std::ranges::input_range R>
            requires layout_or_array<std::ranges::range_value_t<R>>
        void push(R&& arrays)
        {
            arrow_array_stream_private_data& private_data = get_private_data();
 
            // Check if we need to create schema from first array
            if (private_data.schema() == nullptr)
            {
                schema_unique_ptr schema{new ArrowSchema(), arrow_schema_deleter{}};
                copy_schema(*get_arrow_schema(*std::ranges::begin(arrays)), *schema);
                private_data.import_schema(std::move(schema));
            }
 
            // Validate schema compatibility for all arrays
            for (const auto& array : arrays)
            {
                if (!check_compatible_schema(*private_data.schema(), *get_arrow_schema(array)))
                {
                    throw std::runtime_error("Incompatible schema when adding array to ArrowArrayStream");
                }
            }
 
            // Import all arrays
            for (auto&& array : std::forward<R>(arrays))
            {
                ArrowArray extracted_array = extract_arrow_array(std::move(array));
                array_unique_ptr array_ptr{new ArrowArray(), arrow_array_deleter{}};
                swap(*array_ptr, extracted_array);
                private_data.import_array(std::move(array_ptr));
            }
        }

        template <layout_or_array A>
        void push(A&& array)
        {
            push(std::ranges::single_view(std::forward<A>(array)));
        }

        SPARROW_API std::optional<array> pop();
 
    private:
 
        std::variant<ArrowArrayStream*, ArrowArrayStream> m_stream;

        [[nodiscard]] ArrowArrayStream* get_stream_ptr();

        [[nodiscard]] const ArrowArrayStream* get_stream_ptr() const;

        void throw_if_immutable() const;

        [[nodiscard]] SPARROW_API arrow_array_stream_private_data& get_private_data();
    };
}