record_batch.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 mplied.
// See the License for the specific language governing permissions and
// limitations under the License.
 
#pragma once
 
#include <algorithm>
#include <initializer_list>
#include <optional>
#include <ranges>
#include <string>
#include <unordered_map>
#include <vector>
 
#include "sparrow/array.hpp"
#include "sparrow/struct_array.hpp"
#include "sparrow/utils/contracts.hpp"
 
#if defined(__cpp_lib_format)
#    include "sparrow/utils/format.hpp"
#endif
 
namespace sparrow
{
    class record_batch
    {
    public:
 
        using name_type = std::string;
        using size_type = std::size_t;
        using initializer_type = std::initializer_list<std::pair<name_type, array>>;
 
        using name_range = std::ranges::ref_view<const std::vector<name_type>>;

        record_batch() = default;

        template <
            std::ranges::input_range NR,
            std::ranges::input_range CR,
            input_metadata_container METADATA_RANGE = std::vector<metadata_pair>>
            requires(
                std::convertible_to<std::ranges::range_value_t<NR>, std::string>
                and std::same_as<std::ranges::range_value_t<CR>, array>
            )
        constexpr record_batch(
            NR&& names,
            CR&& columns,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        );

        template <std::ranges::input_range CR, input_metadata_container METADATA_RANGE = std::vector<metadata_pair>>
            requires std::same_as<std::ranges::range_value_t<CR>, array>
        record_batch(
            CR&& columns,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        );

        SPARROW_API record_batch(initializer_type init);

        SPARROW_API record_batch(ArrowArray&& array, ArrowSchema&& schema);

        SPARROW_API record_batch(ArrowArray&& array, ArrowSchema* schema);

        SPARROW_API record_batch(ArrowArray&& array, const ArrowSchema* schema);

        SPARROW_API record_batch(ArrowArray* array, ArrowSchema* schema);

        SPARROW_API record_batch(const ArrowArray* array, const ArrowSchema* schema);

        SPARROW_API record_batch(struct_array&& ar);

        SPARROW_API record_batch(const record_batch& other);

        SPARROW_API record_batch& operator=(const record_batch& other);
 
        record_batch(record_batch&&) noexcept = default;
        record_batch& operator=(record_batch&&) = default;

        SPARROW_API size_type nb_columns() const;

        SPARROW_API size_type nb_rows() const;

        SPARROW_API bool contains_column(const name_type& key) const;

        SPARROW_API const name_type& get_column_name(size_type index) const;

        SPARROW_API const array& get_column(const name_type& key) const;

        SPARROW_API array& get_column(const name_type& key);

        SPARROW_API const array& get_column(size_type index) const;

        SPARROW_API array& get_column(size_type index);

        SPARROW_API const std::optional<name_type>& name() const;

        SPARROW_API name_range names() const;

        auto columns() const
        {
            return std::views::iota(size_type{0}, nb_columns())
                   | std::views::transform(
                       [this](size_type i) -> const array&
                       {
                           return get_column(i);
                       }
                   );
        }

        SPARROW_API struct_array extract_struct_array();

        SPARROW_API void add_column(name_type name, array column);

        SPARROW_API void add_column(array column);

        SPARROW_API void add_column_reference(name_type name, array& column);

        SPARROW_API void add_column_reference(array& column);
 
    private:
 
        template <class AS>
        void init(ArrowArray&& arr, AS* sch);
 
        template <class AA, class AS>
        void init(AA* arr, AS* sch);
 
        SPARROW_API void partial_init_from_schema(const ArrowSchema& sch);

        template <class U, class R>
        [[nodiscard]] std::vector<U> to_vector(R&& range) const;

        SPARROW_API void update_array_map_cache() const;

        SPARROW_API void check_consistency() const;
 
        using metadata_type = std::vector<metadata_pair>;
        using array_storage_type = std::variant<array, std::reference_wrapper<array>>;
 
        std::optional<name_type> m_name = std::nullopt;          
        std::optional<metadata_type> m_metadata = std::nullopt;  
        std::vector<name_type> m_name_list;                      
        std::vector<array_storage_type> m_array_list;            
        mutable std::unordered_map<name_type, array*> m_array_map;  
        mutable bool m_dirty_map = true;                            

        SPARROW_API static array* get_array_ptr(array_storage_type& storage);
        SPARROW_API static const array* get_array_ptr(const array_storage_type& storage);
    };

    SPARROW_API
    bool operator==(const record_batch& lhs, const record_batch& rhs);
 
    SPARROW_API std::pair<ArrowArray, ArrowSchema> extract_arrow_structures(sparrow::record_batch&& rb);
 
    /*******************************
     * record_batch implementation *
     *******************************/
 
    template <std::ranges::input_range NR, std::ranges::input_range CR, input_metadata_container METADATA_RANGE>
        requires(std::convertible_to<std::ranges::range_value_t<NR>, std::string>
                 and std::same_as<std::ranges::range_value_t<CR>, array>)
    constexpr record_batch::record_batch(
        NR&& names,
        CR&& columns,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    )
        : m_name(name)
        , m_metadata(std::move(metadata))
        , m_name_list(to_vector<name_type>(std::forward<NR>(names)))
    {
        m_array_list.reserve(std::ranges::size(columns));
        for (auto& col : columns)
        {
            if constexpr (std::is_lvalue_reference_v<CR>)
            {
                m_array_list.emplace_back(col);
            }
            else
            {
                m_array_list.emplace_back(std::move(col));
            }
        }
        update_array_map_cache();
    }
 
    namespace detail
    {
        inline std::vector<record_batch::name_type> get_names(const std::vector<array>& array_list)
        {
            const auto names = array_list
                               | std::views::transform(
                                   [](const array& ar)
                                   {
                                       return ar.name().value();
                                   }
                               );
            return {names.begin(), names.end()};
        }
    }
 
    template <std::ranges::input_range CR, input_metadata_container METADATA_RANGE>
        requires std::same_as<std::ranges::range_value_t<CR>, array>
    record_batch::record_batch(CR&& columns, std::optional<std::string_view> name, std::optional<METADATA_RANGE> metadata)
        : m_name(name)
        , m_metadata(std::move(metadata))
        , m_name_list(detail::get_names(columns))
    {
        m_array_list.reserve(std::ranges::size(columns));
        for (auto& col : columns)
        {
            if constexpr (std::is_lvalue_reference_v<CR>)
            {
                m_array_list.emplace_back(col);
            }
            else
            {
                m_array_list.emplace_back(std::move(col));
            }
        }
        update_array_map_cache();
    }
 
    template <class AS>
    void record_batch::init(ArrowArray&& arr, AS* sch)
    {
        partial_init_from_schema(*sch);
        std::size_t column_size = m_name_list.capacity();
        for (std::size_t i = 0; i < column_size; ++i)
        {
            m_name_list.emplace_back(sch->children[i]->name);
            m_array_list.emplace_back(array(std::move(*(arr.children[i])), sch->children[i]));
            *(arr.children[i]) = make_empty_arrow_array();
        }
        arr.release(&arr);
        update_array_map_cache();
    }
 
    template <class AA, class AS>
    void record_batch::init(AA* arr, AS* sch)
    {
        partial_init_from_schema(*sch);
        std::size_t column_size = m_name_list.capacity();
        for (std::size_t i = 0; i < column_size; ++i)
        {
            m_name_list.emplace_back(sch->children[i]->name);
            m_array_list.emplace_back(array(arr->children[i], sch->children[i]));
        }
        update_array_map_cache();
    }
 
    template <class U, class R>
    std::vector<U> record_batch::to_vector(R&& range) const
    {
        std::vector<U> v;
        if constexpr (std::ranges::sized_range<decltype(range)>)
        {
            v.reserve(std::ranges::size(range));
        }
        if constexpr (std::is_lvalue_reference_v<R>)
        {
            std::ranges::copy(range, std::back_inserter(v));
        }
        else
        {
            std::ranges::move(range, std::back_inserter(v));
        }
        return v;
    }
}
 
#if defined(__cpp_lib_format)
template <>
struct std::formatter<sparrow::record_batch>
{
    constexpr auto parse(std::format_parse_context& ctx)
    {
        return ctx.begin();  // Simple implementation
    }
 
    auto format(const sparrow::record_batch& rb, std::format_context& ctx) const
    {
        auto columns_view = rb.columns();
        std::vector<std::vector<sparrow::array_traits::const_reference>> values_by_columns;
        values_by_columns.reserve(rb.nb_columns());
 
        for (const auto& ar : columns_view)
        {
            std::vector<sparrow::array_traits::const_reference> column_values;
            column_values.reserve(rb.nb_rows());
            for (std::size_t i = 0; i < rb.nb_rows(); ++i)
            {
                column_values.push_back(ar[i]);
            }
            values_by_columns.push_back(std::move(column_values));
        }
 
        sparrow::to_table_with_columns(ctx.out(), rb.names(), values_by_columns);
        return ctx.out();
    }
};
 
namespace sparrow
{
    inline std::ostream& operator<<(std::ostream& os, const record_batch& value)
    {
        os << std::format("{}", value);
        return os;
    }
}
 
#endif