arrow_array_schema_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 <cstdint>
#include <iterator>
#include <optional>
#include <ranges>
#include <string>
#include <string_view>
#include <unordered_set>
 
#include "sparrow/arrow_interface/arrow_array/private_data.hpp"
#include "sparrow/arrow_interface/arrow_array_schema_info_utils.hpp"
#include "sparrow/arrow_interface/arrow_schema/private_data.hpp"
#include "sparrow/buffer/buffer_view.hpp"
#include "sparrow/buffer/dynamic_bitset/dynamic_bitset_view.hpp"
#include "sparrow/buffer/dynamic_bitset/non_owning_dynamic_bitset.hpp"
#include "sparrow/c_interface.hpp"
#include "sparrow/config/config.hpp"
#include "sparrow/types/data_type.hpp"
#include "sparrow/utils/metadata.hpp"
#include "sparrow/utils/ranges.hpp"
 
namespace sparrow
{
    class arrow_proxy_exception : public std::runtime_error
    {
    public:

        explicit arrow_proxy_exception(const std::string& message)
            : std::runtime_error(message)
        {
        }
    };
 
    struct arrow_array_and_schema_pointers
    {
        ArrowArray* array;
        ArrowSchema* schema;
    };
 
    struct arrow_array_and_schema
    {
        ArrowArray array;
        ArrowSchema schema;
    };

    class arrow_proxy
    {
    public:
 
        using bitmap_type = non_owning_dynamic_bitset<uint8_t>;
        using const_bitmap_type = dynamic_bitset_view<const uint8_t>;

        SPARROW_API explicit arrow_proxy(ArrowArray&& array, ArrowSchema&& schema);

        SPARROW_API explicit arrow_proxy(ArrowArray&& array, ArrowSchema* schema);

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

        SPARROW_API explicit arrow_proxy(ArrowArray* array, ArrowSchema* schema);

        SPARROW_API explicit arrow_proxy(const ArrowArray* array, const ArrowSchema* schema);

        SPARROW_API arrow_proxy(const arrow_proxy& other);

        SPARROW_API arrow_proxy& operator=(const arrow_proxy& other);

        SPARROW_API arrow_proxy(arrow_proxy&& other) noexcept;

        SPARROW_API arrow_proxy& operator=(arrow_proxy&& other) noexcept;

        SPARROW_API ~arrow_proxy();

        [[nodiscard]] SPARROW_API const std::string_view format() const;

        SPARROW_API void set_format(const std::string_view format);

        [[nodiscard]] SPARROW_API enum data_type data_type() const;

        void SPARROW_API set_data_type(enum data_type data_type);

        [[nodiscard]] SPARROW_API std::optional<std::string_view> name() const;

        SPARROW_API void set_name(std::optional<std::string_view> name);

        [[nodiscard]] SPARROW_API std::optional<key_value_view> metadata() const;

        template <input_metadata_container R>
        void set_metadata(std::optional<R> metadata)
        {
            if (!schema_created_with_sparrow())
            {
                throw arrow_proxy_exception("Cannot set metadata on non-sparrow created ArrowArray");
            }
            auto private_data = get_schema_private_data();
            if (!metadata.has_value())
            {
                private_data->metadata() = std::nullopt;
            }
            else
            {
                private_data->metadata() = get_metadata_from_key_values(*metadata);
            }
            schema().metadata = private_data->metadata_ptr();
        }

        [[nodiscard]] SPARROW_API std::unordered_set<ArrowFlag> flags() const;

        SPARROW_API void set_flags(const std::unordered_set<ArrowFlag>& flags);

        [[nodiscard]] SPARROW_API size_t length() const;

        SPARROW_API void set_length(size_t length);

        [[nodiscard]] SPARROW_API int64_t null_count() const;

        SPARROW_API void set_null_count(int64_t null_count);

        [[nodiscard]] SPARROW_API size_t offset() const;

        SPARROW_API void set_offset(size_t offset);

        [[nodiscard]] SPARROW_API size_t n_buffers() const;

        SPARROW_API void set_n_buffers(size_t n_buffers);

        [[nodiscard]] SPARROW_API size_t n_children() const;

        [[nodiscard]] SPARROW_API const std::vector<sparrow::buffer_view<uint8_t>>& buffers() const;

        [[nodiscard]] SPARROW_API std::vector<sparrow::buffer_view<uint8_t>>& buffers();

        SPARROW_API void set_buffer(size_t index, const buffer_view<uint8_t>& buffer);

        SPARROW_API void set_buffer(size_t index, buffer<uint8_t>&& buffer);

        SPARROW_API void resize_bitmap(size_t new_size, bool value = true);

        SPARROW_API size_t insert_bitmap(size_t index, bool value, size_t count = 1);

        template <std::ranges::input_range R>
        size_t insert_bitmap(size_t index, const R& range);

        SPARROW_API size_t erase_bitmap(size_t index, size_t count = 1);

        SPARROW_API void push_back_bitmap(bool value);

        SPARROW_API void pop_back_bitmap();

        template <std::ranges::input_range R>
            requires std::same_as<std::ranges::range_value_t<R>, arrow_array_and_schema_pointers>
        void add_children(const R& arrow_array_and_schema_pointers);

        template <std::ranges::input_range R>
            requires std::same_as<std::ranges::range_value_t<R>, arrow_array_and_schema>
        void add_children(R&& arrow_array_and_schemas);

        SPARROW_API void add_child(ArrowArray* array, ArrowSchema* schema);

        SPARROW_API void add_child(const ArrowArray* array, const ArrowSchema* schema);

        SPARROW_API void add_child(ArrowArray&& array, ArrowSchema&& schema);

        SPARROW_API void pop_children(size_t n);

        SPARROW_API void set_child(size_t index, ArrowArray* array, ArrowSchema* schema);

        SPARROW_API void set_child(size_t index, const ArrowArray* array, const ArrowSchema* schema);

        SPARROW_API void set_child(size_t index, ArrowArray&& array, ArrowSchema&& schema);
 

        [[nodiscard]] SPARROW_API const std::vector<arrow_proxy>& children() const;

        [[nodiscard]] SPARROW_API std::vector<arrow_proxy>& children();

        [[nodiscard]] SPARROW_API const std::unique_ptr<arrow_proxy>& dictionary() const;

        [[nodiscard]] SPARROW_API std::unique_ptr<arrow_proxy>& dictionary();

        SPARROW_API void set_dictionary(ArrowArray* array, ArrowSchema* schema);

        SPARROW_API void set_dictionary(const ArrowArray* array, const ArrowSchema* schema);

        SPARROW_API void set_dictionary(ArrowArray&& array_dictionary, ArrowSchema&& schema_dictionary);

        [[nodiscard]] SPARROW_API bool is_created_with_sparrow() const;
 
        [[nodiscard]] SPARROW_API void* private_data() const;

        [[nodiscard]] SPARROW_API arrow_proxy view() const;

        [[nodiscard]] SPARROW_API bool is_view() const noexcept;

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

        [[nodiscard]] SPARROW_API ArrowArray extract_array();

        [[nodiscard]] SPARROW_API ArrowArray& array();

        [[nodiscard]] SPARROW_API const ArrowArray& array() const;

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

        [[nodiscard]] SPARROW_API ArrowSchema extract_schema();

        [[nodiscard]] SPARROW_API ArrowSchema& schema();

        [[nodiscard]] SPARROW_API const ArrowSchema& schema() const;
 
        [[nodiscard]] SPARROW_API arrow_schema_private_data* get_schema_private_data();
        [[nodiscard]] SPARROW_API arrow_array_private_data* get_array_private_data();

        [[nodiscard]] SPARROW_API arrow_proxy slice(size_t start, size_t end) const;

        [[nodiscard]] SPARROW_API arrow_proxy slice_view(size_t start, size_t end) const;

        SPARROW_API void slice_inplace(size_t start, size_t end);

        SPARROW_API void update_buffers();

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

        [[nodiscard]] SPARROW_API bool is_schema_const() const;
 
        [[nodiscard]] SPARROW_API std::optional<bitmap_type>& bitmap();
 
        [[nodiscard]] SPARROW_API const std::optional<const_bitmap_type>& const_bitmap() const;
 
    private:
 
        std::variant<ArrowArray*, ArrowArray> m_array;
        std::variant<ArrowSchema*, ArrowSchema> m_schema;
        std::vector<sparrow::buffer_view<uint8_t>> m_buffers;
        std::vector<arrow_proxy> m_children;
        std::unique_ptr<arrow_proxy> m_dictionary;
        bool m_array_is_immutable = false;
        bool m_schema_is_immutable = false;
        bool m_is_dictionary_immutable = false;
        std::vector<bool> m_children_array_immutable;
        std::vector<bool> m_children_schema_immutable;
        std::optional<bitmap_type> m_null_bitmap;
        std::optional<const_bitmap_type> m_const_bitmap;
 
        struct impl_tag
        {
        };
 
        // Build an empty proxy. Convenient for resizing vector of children
        arrow_proxy();
 
        template <typename AA, typename AS>
            requires std::same_as<std::remove_const_t<std::remove_pointer_t<std::remove_cvref_t<AA>>>, ArrowArray>
                     && std::same_as<std::remove_const_t<std::remove_pointer_t<std::remove_cvref_t<AS>>>, ArrowSchema>
        arrow_proxy(AA&& array, AS&& schema, impl_tag);
 
        [[nodiscard]] bool empty() const;
        SPARROW_API void resize_children(size_t children_count);
 
        void update_children();
        void update_dictionary();
        void update_null_count();
        void reset();
        void remove_dictionary();
        void remove_child(size_t index);
        void create_bitmap_view(std::optional<size_t> null_count = std::nullopt);
 
        [[nodiscard]] bool array_created_with_sparrow() const;
        [[nodiscard]] SPARROW_API bool schema_created_with_sparrow() const;
 
        void validate_array_and_schema() const;
 
        [[nodiscard]] bool is_arrow_array_valid() const;
        [[nodiscard]] bool is_arrow_schema_valid() const;
        [[nodiscard]] bool is_proxy_valid() const;
 
        [[nodiscard]] size_t get_null_count() const;
 
        [[nodiscard]] ArrowArray& array_without_sanitize();
        [[nodiscard]] const ArrowArray& array_without_sanitize() const;
 
        [[nodiscard]] ArrowSchema& schema_without_sanitize();
        [[nodiscard]] const ArrowSchema& schema_without_sanitize() const;

        void sanitize_schema();
 
        void swap(arrow_proxy& other) noexcept;
 
        template <const char* function_name, bool check_array_is_mutable, bool check_schema_is_mutable>
        void throw_if_immutable() const;
    };
 
    template <std::ranges::input_range R>
        requires std::same_as<std::ranges::range_value_t<R>, arrow_array_and_schema_pointers>
    void arrow_proxy::add_children(const R& arrow_array_and_schema_pointers)
    {
        static constexpr const char function_name[] = "add_children";
        throw_if_immutable<function_name, true, true>();
        const size_t add_children_count = std::ranges::size(arrow_array_and_schema_pointers);
        const size_t original_children_count = n_children();
        const size_t new_children_count = original_children_count + add_children_count;
 
        resize_children(new_children_count);
        for (size_t i = 0; i < add_children_count; ++i)
        {
            set_child(
                i + original_children_count,
                arrow_array_and_schema_pointers[i].array,
                arrow_array_and_schema_pointers[i].schema
            );
        }
    }
 
    template <std::ranges::input_range R>
        requires std::same_as<std::ranges::range_value_t<R>, arrow_array_and_schema>
    void arrow_proxy::add_children(R&& arrow_arrays_and_schemas)
    {
        if (!is_created_with_sparrow())
        {
            throw arrow_proxy_exception("Cannot set n_buffers on non-sparrow created ArrowArray or ArrowSchema");
        }
 
        const size_t add_children_count = std::ranges::size(arrow_arrays_and_schemas);
        const size_t original_children_count = n_children();
        const size_t new_children_count = original_children_count + add_children_count;
 
        resize_children(new_children_count);
        for (size_t i = 0; i < add_children_count; ++i)
        {
            set_child(
                i + original_children_count,
                std::move(arrow_arrays_and_schemas[i].array),
                std::move(arrow_arrays_and_schemas[i].schema)
            );
        }
    }
 
    template <std::ranges::input_range R>
    inline size_t arrow_proxy::insert_bitmap(size_t index, const R& range)
    {
        static constexpr const char function_name[] = "insert_bitmap";
        throw_if_immutable<function_name, true, false>();
        SPARROW_ASSERT_TRUE(m_null_bitmap.has_value())
        const auto it = m_null_bitmap->insert(
            sparrow::next(m_null_bitmap->cbegin(), index),
            range.begin(),
            range.end()
        );
        set_null_count(static_cast<int64_t>(m_null_bitmap->null_count()));
        m_const_bitmap = const_bitmap_type(
            m_null_bitmap->data(),
            m_null_bitmap->size(),
            static_cast<size_t>(m_null_bitmap->offset()),
            static_cast<size_t>(m_null_bitmap->null_count())
        );
        update_buffers();
        return static_cast<size_t>(std::distance(m_null_bitmap->begin(), it));
    }
 
    template <typename AA, typename AS>
        requires std::same_as<std::remove_const_t<std::remove_pointer_t<std::remove_cvref_t<AA>>>, ArrowArray>
                 && std::same_as<std::remove_const_t<std::remove_pointer_t<std::remove_cvref_t<AS>>>, ArrowSchema>
    arrow_proxy::arrow_proxy(AA&& array, AS&& schema, impl_tag)
    {
        if constexpr (std::is_const_v<std::remove_pointer_t<std::remove_reference_t<AA>>>)
        {
            m_array_is_immutable = true;
            m_array = const_cast<ArrowArray*>(array);
        }
        else
        {
            m_array = std::forward<AA>(array);
        }
 
        if constexpr (std::is_const_v<std::remove_pointer_t<std::remove_reference_t<AS>>>)
        {
            m_schema_is_immutable = true;
            m_schema = const_cast<ArrowSchema*>(schema);
        }
        else
        {
            m_schema = std::forward<AS>(schema);
        }
 
        if constexpr (std::is_rvalue_reference_v<AA&&>)
        {
            array = {};
        }
        else if constexpr (std::is_pointer_v<std::remove_cvref_t<AA>>)
        {
            SPARROW_ASSERT_TRUE(array != nullptr);
        }
 
        if constexpr (std::is_rvalue_reference_v<AS&&>)
        {
            schema = {};
        }
        else if constexpr (std::is_pointer_v<std::remove_cvref_t<AS>>)
        {
            SPARROW_ASSERT_TRUE(schema != nullptr);
        }
 
        m_children_array_immutable = std::vector<bool>(n_children(), m_array_is_immutable);
        m_children_schema_immutable = std::vector<bool>(n_children(), m_schema_is_immutable);
        validate_array_and_schema();
        update_buffers();
        update_children();
        update_dictionary();
        create_bitmap_view();
    }
 
    template <const char* function_name, bool check_array_is_mutable, bool check_schema_is_mutable>
    void arrow_proxy::throw_if_immutable() const
    {
        static const std::string cannot_call = "Cannot call ";
        if (!is_created_with_sparrow())
        {
            auto error_message = cannot_call + std::string(function_name)
                                 + " on non-sparrow created ArrowArray or ArrowSchema";
            throw arrow_proxy_exception(error_message);
        }
        if constexpr (check_array_is_mutable || check_schema_is_mutable)
        {
            if (m_array_is_immutable || m_schema_is_immutable)
            {
                {
                    std::string error_message = cannot_call + std::string(function_name);
                    if constexpr (check_array_is_mutable && !check_schema_is_mutable)
                    {
                        if (m_array_is_immutable)
                        {
                            error_message += " on an immutable ArrowArray. You may have passed a const ArrowArray* at the creation.";
                        }
                    }
                    else if constexpr (check_schema_is_mutable && !check_array_is_mutable)
                    {
                        if (m_schema_is_immutable)
                        {
                            error_message += " on an immutable ArrowSchema. You may have passed a const ArrowSchema* at the creation.";
                        }
                    }
                    else if constexpr (check_array_is_mutable && check_schema_is_mutable)
                    {
                        if (m_array_is_immutable && m_schema_is_immutable)
                        {
                            error_message += " on an immutable ArrowArray and ArrowSchema. You may have passed const ArrowArray* and const ArrowSchema* at the creation.";
                        }
                    }
                    throw arrow_proxy_exception(error_message);
                }
            }
        }
    }
}
 
#if defined(__cpp_lib_format)
 
template <>
struct std::formatter<sparrow::buffer_view<uint8_t>>
{
private:
 
    char delimiter = ' ';
    static constexpr std::string_view opening = "[";
    static constexpr std::string_view closing = "]";
 
public:
 
    constexpr auto parse(std::format_parse_context& ctx)
    {
        auto it = ctx.begin();
        auto end = ctx.end();
 
        // Parse optional delimiter
        if (it != end && *it != '}')
        {
            delimiter = *it++;
        }
 
        if (it != end && *it != '}')
        {
            throw std::format_error("Invalid format specifier for range");
        }
 
        return it;
    }
 
    auto format(const sparrow::buffer_view<uint8_t>& range, std::format_context& ctx) const
    {
        auto out = ctx.out();
 
        // Write opening bracket
        out = sparrow::ranges::copy(opening, out).out;
 
        // Write range elements
        bool first = true;
        for (const auto& elem : range)
        {
            if (!first)
            {
                *out++ = delimiter;
            }
            out = std::format_to(out, "{}", elem);
            first = false;
        }
 
        // Write closing bracket
        out = sparrow::ranges::copy(closing, out).out;
 
        return out;
    }
};
 
namespace sparrow
{
    inline std::ostream& operator<<(std::ostream& os, const buffer_view<uint8_t>& value)
    {
        os << std::format("{}", value);
        return os;
    }
}
 
template <>
struct std::formatter<sparrow::arrow_proxy>
{
    constexpr auto parse(std::format_parse_context& ctx)
    {
        return ctx.begin();  // Simple implementation
    }
 
    auto format(const sparrow::arrow_proxy& obj, std::format_context& ctx) const
    {
        std::string buffers_description_str;
        for (size_t i = 0; i < obj.n_buffers(); ++i)
        {
            std::format_to(
                std::back_inserter(buffers_description_str),
                "<{}[{} b]{}",
                "uint8_t",
                obj.buffers()[i].size() * sizeof(uint8_t),
                obj.buffers()[i]
            );
        }
 
        std::string children_str;
        for (const auto& child : obj.children())
        {
            std::format_to(std::back_inserter(children_str), "{}\n", child);
        }
 
        const std::string dictionary_str = obj.dictionary() ? std::format("{}", *obj.dictionary()) : "nullptr";
        const std::string metadata_str = obj.metadata() ? std::format("{}", *obj.metadata()) : "nullptr";
 
        return std::format_to(
            ctx.out(),
            "arrow_proxy\n- format: {}\n- name; {}\n- metadata: {}\n- data_type: {}\n- null_count:{}\n- length: {}\n- offset: {}\n- n_buffers: {}\n- buffers:\n{}\n- n_children: {}\n-children: {}\n- dictionary: {}",
            obj.format(),
            obj.name().value_or(""),
            metadata_str,
            obj.data_type(),
            obj.null_count(),
            obj.length(),
            obj.offset(),
            obj.n_buffers(),
            buffers_description_str,
            obj.n_children(),
            children_str,
            dictionary_str
        );
    }
};
 
namespace sparrow
{
    inline std::ostream& operator<<(std::ostream& os, const arrow_proxy& value)
    {
        os << std::format("{}", value);
        return os;
    }
}
 
#endif