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/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:

        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, ArrowSchema* schema);
 
        // Copy constructors
        SPARROW_API arrow_proxy(const arrow_proxy&);
        SPARROW_API arrow_proxy& operator=(const arrow_proxy&);
 
        // Move constructors
        SPARROW_API arrow_proxy(arrow_proxy&&) noexcept;
        SPARROW_API arrow_proxy& operator=(arrow_proxy&&) 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(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, 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(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 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 update_buffers();
 
    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;
 
        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_pointer_t<std::remove_cvref_t<AA>>, ArrowArray>
                     && std::same_as<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);
 
        [[nodiscard]] SPARROW_API non_owning_dynamic_bitset<uint8_t> get_non_owning_dynamic_bitset();
 
        void update_children();
        void update_dictionary();
        void update_null_count();
        void reset();
 
        [[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 <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)
    {
        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_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)
    {
        if (!is_created_with_sparrow())
        {
            throw arrow_proxy_exception("Cannot modify the bitmap on non-sparrow created ArrowArray");
        }
        SPARROW_ASSERT_TRUE(has_bitmap(data_type()))
        auto bitmap = get_non_owning_dynamic_bitset();
        const auto it = bitmap.insert(sparrow::next(bitmap.cbegin(), index), range.begin(), range.end());
        return static_cast<size_t>(std::distance(bitmap.begin(), it));
    }
}
 
#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;
    }
};
 
inline std::ostream& operator<<(std::ostream& os, const sparrow::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";
 
        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(""),
            obj.metadata().value_or(""),
            obj.data_type(),
            obj.null_count(),
            obj.length(),
            obj.offset(),
            obj.n_buffers(),
            buffers_description_str,
            obj.n_children(),
            children_str,
            dictionary_str
        );
    }
};
 
inline std::ostream& operator<<(std::ostream& os, const sparrow::arrow_proxy& value)
{
    os << std::format("{}", value);
    return os;
}
 
#endif