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 <ranges>
#include <string>
#include <unordered_map>
#include <vector>
 
#include "sparrow/array.hpp"
#include "sparrow/layout/struct_layout/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>>;
        using column_range = std::ranges::ref_view<const std::vector<array>>;

        template <std::ranges::input_range NR, std::ranges::input_range CR>
            requires(
                std::convertible_to<std::ranges::range_value_t<NR>, std::string>
                and std::same_as<std::ranges::range_value_t<CR>, array>
            )
        record_batch(NR&& names, CR&& columns);
 
        /*
         * Constructs a @ref record_batch from a range of arrays. Each array
         * must have a name: if \c arr is an array, then \c arr.name(), must
         * not return an empty string.
         *
         * @param comumns An input range of arrays
         */
        template <std::ranges::input_range CR>
            requires std::same_as<std::ranges::range_value_t<CR>, array>
        record_batch(CR&& columns);

        SPARROW_API record_batch(initializer_type init);

        SPARROW_API record_batch(struct_array&& ar);
 
        SPARROW_API record_batch(const record_batch&);
        SPARROW_API record_batch& operator=(const record_batch&);
 
        record_batch(record_batch&&) = 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 const array& get_column(size_type index) const;

        SPARROW_API name_range names() const;

        SPARROW_API column_range columns() const;

        SPARROW_API struct_array extract_struct_array();

        SPARROW_API void add_column(name_type name, array column);

        SPARROW_API void add_column(array column);
 
    private:
 
        template <class U, class R>
        [[nodiscard]] std::vector<U> to_vector(R&& range) const;
 
        SPARROW_API void update_array_map_cache() const;
 
        [[nodiscard]] SPARROW_API bool check_consistency() const;
 
        std::vector<name_type> m_name_list;
        std::vector<array> m_array_list;
        mutable std::unordered_map<name_type, const array*> m_array_map;
        mutable bool m_dirty_map = true;
    };

    SPARROW_API
    bool operator==(const record_batch& lhs, const record_batch& rhs);
 
    /*******************************
     * record_batch implementation *
     *******************************/
 
    template <std::ranges::input_range NR, std::ranges::input_range CR>
        requires(std::convertible_to<std::ranges::range_value_t<NR>, std::string>
                 and std::same_as<std::ranges::range_value_t<CR>, array>)
    record_batch::record_batch(NR&& names, CR&& columns)
        : m_name_list(to_vector<name_type>(std::move(names)))
        , m_array_list(to_vector<array>(std::move(columns)))
    {
        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>
        requires std::same_as<std::ranges::range_value_t<CR>, array>
    record_batch::record_batch(CR&& columns)
        : m_name_list(detail::get_names(columns))
        , m_array_list(to_vector<array>(std::move(columns)))
    {
        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));
        }
        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
    {
        const auto values_by_columns = rb.columns()
                                       | std::views::transform(
                                           [&rb](const auto& ar)
                                           {
                                               return std::views::iota(0u, rb.nb_rows())
                                                      | std::views::transform(
                                                          [&ar](const auto i)
                                                          {
                                                              return ar[i];
                                                          }
                                                      );
                                           }
                                       );
 
        sparrow::to_table_with_columns(ctx.out(), rb.names(), values_by_columns);
        return ctx.out();
    }
};
 
inline std::ostream& operator<<(std::ostream& os, const sparrow::record_batch& value)
{
    os << std::format("{}", value);
    return os;
}
 
#endif