dictionary_encoded_array.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 <optional>
 
#include "sparrow/array_api.hpp"
#include "sparrow/array_factory.hpp"
#include "sparrow/arrow_array_schema_proxy.hpp"
#include "sparrow/buffer/dynamic_bitset/dynamic_bitset.hpp"
#include "sparrow/c_interface.hpp"
#include "sparrow/layout/array_access.hpp"
#include "sparrow/layout/array_helper.hpp"
#include "sparrow/layout/array_wrapper.hpp"
#include "sparrow/layout/nested_value_types.hpp"
#include "sparrow/layout/primitive_layout/primitive_array.hpp"
#include "sparrow/types/data_traits.hpp"
#include "sparrow/utils/contracts.hpp"
#include "sparrow/utils/functor_index_iterator.hpp"
#include "sparrow/utils/memory.hpp"
#include "sparrow/utils/metadata.hpp"
#include "sparrow/utils/mp_utils.hpp"
 
namespace sparrow
{
    template <class Layout, bool is_const>
    class layout_element_functor
    {
    public:
 
        using layout_type = Layout;
        using storage_type = std::conditional_t<is_const, const layout_type*, layout_type>;
        using return_type = std::
            conditional_t<is_const, typename layout_type::const_reference, typename layout_type::reference>;
 
        constexpr layout_element_functor() = default;
 
        constexpr explicit layout_element_functor(storage_type layout_)
            : p_layout(layout_)
        {
        }
 
        [[nodiscard]] return_type operator()(std::size_t i) const
        {
            return p_layout->operator[](i);
        }
 
    private:
 
        storage_type p_layout;
    };
 
    template <std::integral IT>
    class dictionary_encoded_array;
 
    namespace detail
    {
        template <class T>
        struct get_data_type_from_array;
 
        template <std::integral IT>
        struct get_data_type_from_array<sparrow::dictionary_encoded_array<IT>>
        {
            [[nodiscard]] static constexpr sparrow::data_type get()
            {
                return arrow_traits<typename primitive_array<IT>::inner_value_type>::type_id;
            }
        };
 
        template <std::integral IT>
        struct is_dictionary_encoded_array<sparrow::dictionary_encoded_array<IT>>
        {
            [[nodiscard]] static constexpr bool get()
            {
                return true;
            }
        };
    }

    template <class T>
    constexpr bool is_dictionary_encoded_array_v = detail::is_dictionary_encoded_array<T>::get();
 
    template <std::integral IT>
    class dictionary_encoded_array
    {
    public:
 
        using self_type = dictionary_encoded_array<IT>;
        using size_type = std::size_t;
        using difference_type = std::ptrdiff_t;
 
        using inner_value_type = array_traits::inner_value_type;
 
        using value_type = array_traits::value_type;
        using reference = array_traits::const_reference;
        using const_reference = array_traits::const_reference;
 
        using functor_type = layout_element_functor<self_type, true>;
        using const_functor_type = layout_element_functor<self_type, true>;
 
        using iterator = functor_index_iterator<functor_type>;
        using const_iterator = functor_index_iterator<const_functor_type>;
 
        using keys_buffer_type = u8_buffer<IT>;
 
        explicit dictionary_encoded_array(arrow_proxy);
 
        dictionary_encoded_array(const self_type&);
        self_type& operator=(const self_type&);
 
        dictionary_encoded_array(self_type&&);
        self_type& operator=(self_type&&);
 
        [[nodiscard]] std::optional<std::string_view> name() const;
        [[nodiscard]] std::optional<key_value_view> metadata() const;
 
        [[nodiscard]] size_type size() const;
        [[nodiscard]] bool empty() const;
 
        [[nodiscard]] const_reference operator[](size_type i) const;
 
        [[nodiscard]] iterator begin();
        [[nodiscard]] iterator end();
 
        [[nodiscard]] const_iterator begin() const;
        [[nodiscard]] const_iterator end() const;
 
        [[nodiscard]] const_iterator cbegin() const;
        [[nodiscard]] const_iterator cend() const;
 
        [[nodiscard]] const_reference front() const;
        [[nodiscard]] const_reference back() const;
 
        template <class... Args>
            requires(mpl::excludes_copy_and_move_ctor_v<dictionary_encoded_array<IT>, Args...>)
        explicit dictionary_encoded_array(Args&&... args)
            : dictionary_encoded_array(create_proxy(std::forward<Args>(args)...))
        {
        }

        [[nodiscard]] self_type slice(size_type start, size_type end) const;

        [[nodiscard]] self_type slice_view(size_type start, size_type end) const;
 
    private:
 
        template <
            validity_bitmap_input R = validity_bitmap,
            input_metadata_container METADATA_RANGE = std::vector<metadata_pair>>
        [[nodiscard]] static auto create_proxy(
            keys_buffer_type&& keys,
            array&& values,
            R&& bitmaps,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        ) -> arrow_proxy;
 
        template <
            validity_bitmap_input R = validity_bitmap,
            input_metadata_container METADATA_RANGE = std::vector<metadata_pair>>
        [[nodiscard]] static auto create_proxy(
            keys_buffer_type&& keys,
            array&& values,
            bool nullable = true,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        ) -> arrow_proxy;
 
        template <input_metadata_container METADATA_RANGE = std::vector<metadata_pair>>
        [[nodiscard]] static auto create_proxy_impl(
            keys_buffer_type&& keys,
            array&& values,
            std::optional<validity_bitmap> validity = std::nullopt,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        ) -> arrow_proxy;
 
        template <
            std::ranges::input_range KEY_RANGE,
            validity_bitmap_input R = validity_bitmap,
            input_metadata_container METADATA_RANGE = std::vector<metadata_pair>>
            requires(
                !std::same_as<KEY_RANGE, keys_buffer_type>
                and std::same_as<IT, std::ranges::range_value_t<KEY_RANGE>>
            )
        [[nodiscard]] static arrow_proxy create_proxy(
            KEY_RANGE&& keys,
            array&& values,
            R&& bitmaps = validity_bitmap{},
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        )
        {
            keys_buffer_type keys_buffer(std::forward<KEY_RANGE>(keys));
            return create_proxy(
                std::move(keys_buffer),
                std::forward<array>(values),
                std::forward<R>(bitmaps),
                std::move(name),
                std::move(metadata)
            );
        }
 
        // range of nullable values
        template <
            std::ranges::input_range NULLABLE_KEY_RANGE,
            input_metadata_container METADATA_RANGE = std::vector<metadata_pair>>
            requires std::is_same_v<std::ranges::range_value_t<NULLABLE_KEY_RANGE>, nullable<IT>>
        static arrow_proxy create_proxy(
            NULLABLE_KEY_RANGE&& nullable_keys,
            array&& values,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        );
 
        using keys_layout = primitive_array<IT>;
        using values_layout = cloning_ptr<array_wrapper>;
 
        [[nodiscard]] const inner_value_type& dummy_inner_value() const;
        [[nodiscard]] const_reference dummy_const_reference() const;
 
        [[nodiscard]] static keys_layout create_keys_layout(arrow_proxy& proxy);
        [[nodiscard]] static values_layout create_values_layout(arrow_proxy& proxy);
 
        [[nodiscard]] arrow_proxy& get_arrow_proxy();
        [[nodiscard]] const arrow_proxy& get_arrow_proxy() const;
 
        arrow_proxy m_proxy;
        keys_layout m_keys_layout;
        values_layout p_values_layout;
 
        friend class detail::array_access;
    };
 
    template <class IT>
    bool operator==(const dictionary_encoded_array<IT>& lhs, const dictionary_encoded_array<IT>& rhs);
 
    /*******************************************
     * dictionary_encoded_array implementation *
     *******************************************/
 
    template <std::integral IT>
    dictionary_encoded_array<IT>::dictionary_encoded_array(arrow_proxy proxy)
        : m_proxy(std::move(proxy))
        , m_keys_layout(create_keys_layout(m_proxy))
        , p_values_layout(create_values_layout(m_proxy))
    {
        SPARROW_ASSERT_TRUE(data_type_is_integer(m_proxy.data_type()));
    }
 
    template <std::integral IT>
    dictionary_encoded_array<IT>::dictionary_encoded_array(const self_type& rhs)
        : m_proxy(rhs.m_proxy)
        , m_keys_layout(create_keys_layout(m_proxy))
        , p_values_layout(create_values_layout(m_proxy))
    {
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::operator=(const self_type& rhs) -> self_type&
    {
        if (this != &rhs)
        {
            m_proxy = rhs.m_proxy;
            m_keys_layout = create_keys_layout(m_proxy);
            p_values_layout = create_values_layout(m_proxy);
        }
        return *this;
    }
 
    template <std::integral IT>
    dictionary_encoded_array<IT>::dictionary_encoded_array(self_type&& rhs)
        : m_proxy(std::move(rhs.m_proxy))
        , m_keys_layout(create_keys_layout(m_proxy))
        , p_values_layout(create_values_layout(m_proxy))
    {
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::operator=(self_type&& rhs) -> self_type&
    {
        if (this != &rhs)
        {
            using std::swap;
            swap(m_proxy, rhs.m_proxy);
            m_keys_layout = create_keys_layout(m_proxy);
            p_values_layout = create_values_layout(m_proxy);
        }
        return *this;
    }
 
    template <std::integral IT>
    template <validity_bitmap_input VBI, input_metadata_container METADATA_RANGE>
    auto dictionary_encoded_array<IT>::create_proxy(
        keys_buffer_type&& keys,
        array&& values,
        VBI&& validity_input,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    ) -> arrow_proxy
    {
        const auto size = keys.size();
        validity_bitmap vbitmap = ensure_validity_bitmap(size, std::forward<VBI>(validity_input));
        return create_proxy_impl(
            std::forward<keys_buffer_type>(keys),
            std::forward<array>(values),
            std::make_optional<validity_bitmap>(std::move(vbitmap)),
            std::move(name),
            std::move(metadata)
        );
    }
 
    template <std::integral IT>
    template <validity_bitmap_input VBI, input_metadata_container METADATA_RANGE>
    auto dictionary_encoded_array<IT>::create_proxy(
        keys_buffer_type&& keys,
        array&& values,
        bool nullable,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    ) -> arrow_proxy
    {
        const auto size = keys.size();
        return create_proxy_impl(
            std::forward<keys_buffer_type>(keys),
            std::forward<array>(values),
            nullable ? std::make_optional<validity_bitmap>(nullptr, size) : std::nullopt,
            std::move(name),
            std::move(metadata)
        );
    }
 
    template <std::integral IT>
    template <input_metadata_container METADATA_RANGE>
    [[nodiscard]] arrow_proxy dictionary_encoded_array<IT>::create_proxy_impl(
        keys_buffer_type&& keys,
        array&& values,
        std::optional<validity_bitmap> validity,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    )
    {
        const auto size = keys.size();
        auto [value_array, value_schema] = extract_arrow_structures(std::move(values));
        static const repeat_view<bool> children_ownership{true, 0};
 
        const std::optional<std::unordered_set<sparrow::ArrowFlag>>
            flags = validity.has_value()
                        ? std::make_optional<std::unordered_set<sparrow::ArrowFlag>>({ArrowFlag::NULLABLE})
                        : std::nullopt;
 
        // create arrow schema and array
        ArrowSchema schema = make_arrow_schema(
            sparrow::data_type_format_of<IT>(),
            std::move(name),                           // name
            std::move(metadata),                       // metadata
            flags,                                     // flags
            nullptr,                                   // children
            children_ownership,                        // children_ownership
            new ArrowSchema(std::move(value_schema)),  // dictionary
            true                                       // dictionary ownership
        );
 
        const size_t null_count = validity.has_value() ? validity->null_count() : 0;
 
        std::vector<buffer<uint8_t>> buffers(2);
        buffers[0] = validity.has_value() ? std::move(*validity).extract_storage()
                                          : buffer<uint8_t>{nullptr, 0};
        buffers[1] = std::move(keys).extract_storage();
        // create arrow array
        ArrowArray arr = make_arrow_array(
            static_cast<std::int64_t>(size),        // length
            static_cast<std::int64_t>(null_count),  // Null count
            0,                                      // offset
            std::move(buffers),
            nullptr,                                 // children
            children_ownership,                      // children_ownership
            new ArrowArray(std::move(value_array)),  // dictionary
            true                                     // dictionary ownership
        );
        return arrow_proxy(std::move(arr), std::move(schema));
    }
 
    template <std::integral IT>
    template <std::ranges::input_range NULLABLE_KEY_RANGE, input_metadata_container METADATA_RANGE>
        requires std::is_same_v<std::ranges::range_value_t<NULLABLE_KEY_RANGE>, nullable<IT>>
    arrow_proxy dictionary_encoded_array<IT>::create_proxy(
        NULLABLE_KEY_RANGE&& nullable_keys,
        array&& values,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    )
    {
        auto keys = nullable_keys
                    | std::views::transform(
                        [](const auto& v)
                        {
                            return v.get();
                        }
                    );
        auto is_non_null = nullable_keys
                           | std::views::transform(
                               [](const auto& v)
                               {
                                   return v.has_value();
                               }
                           );
        return create_proxy(
            std::move(keys),
            std::forward<array>(values),
            std::move(is_non_null),
            std::move(name),
            std::move(metadata)
        );
    }
 
    template <std::integral IT>
    std::optional<std::string_view> dictionary_encoded_array<IT>::name() const
    {
        return m_proxy.name();
    }
 
    template <std::integral IT>
    std::optional<key_value_view> dictionary_encoded_array<IT>::metadata() const
    {
        return m_proxy.metadata();
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::size() const -> size_type
    {
        return m_proxy.length();
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::empty() const -> bool
    {
        return size() == 0;
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::operator[](size_type i) const -> const_reference
    {
        SPARROW_ASSERT_TRUE(i < size());
        const auto index = m_keys_layout[i];
 
        if (index.has_value())
        {
            SPARROW_ASSERT_TRUE(index.value() >= 0);
            return array_element(*p_values_layout, static_cast<std::size_t>(index.value()));
        }
        else
        {
            return dummy_const_reference();
        }
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::begin() -> iterator
    {
        return iterator(functor_type(this), 0u);
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::end() -> iterator
    {
        return iterator(functor_type(this), size());
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::begin() const -> const_iterator
    {
        return cbegin();
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::end() const -> const_iterator
    {
        return cend();
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::cbegin() const -> const_iterator
    {
        return const_iterator(const_functor_type(this), 0u);
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::cend() const -> const_iterator
    {
        return const_iterator(const_functor_type(this), size());
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::front() const -> const_reference
    {
        SPARROW_ASSERT_FALSE(empty());
        return operator[](0);
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::back() const -> const_reference
    {
        SPARROW_ASSERT_FALSE(empty());
        return operator[](size() - 1);
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::dummy_inner_value() const -> const inner_value_type&
    {
        static const inner_value_type instance = array_default_element_value(*p_values_layout);
        return instance;
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::slice(size_type start, size_type end) const -> self_type
    {
        SPARROW_ASSERT_TRUE(start <= end);
        return self_type{get_arrow_proxy().slice(start, end)};
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::slice_view(size_type start, size_type end) const -> self_type
    {
        SPARROW_ASSERT_TRUE(start <= end);
        return self_type{get_arrow_proxy().slice_view(start, end)};
    }
 
    /*template <std::integral IT>
    auto dictionary_encoded_array<IT>::dummy_inner_const_reference() const -> inner_const_reference
    {
        static const inner_const_reference instance =
            std::visit([](const auto& val) -> inner_const_reference { return val; }, dummy_inner_value());
        return instance;
    }*/
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::dummy_const_reference() const -> const_reference
    {
        static const const_reference instance = std::visit(
            [](const auto& val) -> const_reference
            {
                using inner_ref = typename arrow_traits<std::decay_t<decltype(val)>>::const_reference;
                return const_reference{nullable<inner_ref>(inner_ref(val), false)};
            },
            dummy_inner_value()
        );
        return instance;
    }
 
    template <std::integral IT>
    typename dictionary_encoded_array<IT>::values_layout
    dictionary_encoded_array<IT>::create_values_layout(arrow_proxy& proxy)
    {
        const auto& dictionary = proxy.dictionary();
        SPARROW_ASSERT_TRUE(dictionary);
        arrow_proxy ar_dictionary{&(dictionary->array()), &(dictionary->schema())};
        return array_factory(std::move(ar_dictionary));
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::create_keys_layout(arrow_proxy& proxy) -> keys_layout
    {
        return keys_layout{arrow_proxy{&proxy.array(), &proxy.schema()}};
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::get_arrow_proxy() -> arrow_proxy&
    {
        return m_proxy;
    }
 
    template <std::integral IT>
    auto dictionary_encoded_array<IT>::get_arrow_proxy() const -> const arrow_proxy&
    {
        return m_proxy;
    }
 
    template <class IT>
    bool operator==(const dictionary_encoded_array<IT>& lhs, const dictionary_encoded_array<IT>& rhs)
    {
        return std::ranges::equal(lhs, rhs);
    }
}
 
#if defined(__cpp_lib_format)
template <std::integral IT>
struct std::formatter<sparrow::dictionary_encoded_array<IT>>
{
    constexpr auto parse(std::format_parse_context& ctx)
    {
        return ctx.begin();  // Simple implementation
    }
 
    auto format(const sparrow::dictionary_encoded_array<IT>& ar, std::format_context& ctx) const
    {
        std::format_to(ctx.out(), "Dictionary [size={}] <", ar.size());
        std::for_each(
            ar.cbegin(),
            std::prev(ar.cend()),
            [&ctx](const auto& value)
            {
                std::format_to(ctx.out(), "{}, ", value);
            }
        );
        std::format_to(ctx.out(), "{}>", ar.back());
        return ctx.out();
    }
};
 
template <std::integral IT>
std::ostream& operator<<(std::ostream& os, const sparrow::dictionary_encoded_array<IT>& value)
{
    os << std::format("{}", value);
    return os;
}
#endif