variable_size_binary_view_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 implied.
// See the License for the specific language governing permissions and
// limitations under the License.
 
#pragma once
 
#include <cstddef>
#include <ranges>
 
#include "sparrow/arrow_interface/arrow_array.hpp"
#include "sparrow/arrow_interface/arrow_array_schema_proxy.hpp"
#include "sparrow/arrow_interface/arrow_schema.hpp"
#include "sparrow/buffer/dynamic_bitset.hpp"
#include "sparrow/layout/array_access.hpp"
#include "sparrow/layout/array_bitmap_base.hpp"
#include "sparrow/layout/layout_utils.hpp"
#include "sparrow/types/data_traits.hpp"
#include "sparrow/u8_buffer.hpp"
#include "sparrow/utils/functor_index_iterator.hpp"
#include "sparrow/utils/iterator.hpp"
#include "sparrow/utils/mp_utils.hpp"
#include "sparrow/utils/nullable.hpp"
#include "sparrow/utils/ranges.hpp"
#include "sparrow/utils/repeat_container.hpp"
#include "sparrow/utils/sequence_view.hpp"
 
namespace sparrow
{
    template <std::ranges::sized_range T, class CR>
    class variable_size_binary_view_array_impl;

    using string_view_array = variable_size_binary_view_array_impl<
        arrow_traits<std::string>::value_type,
        arrow_traits<std::string>::const_reference>;

    using binary_view_array = variable_size_binary_view_array_impl<
        arrow_traits<std::vector<byte_t>>::value_type,
        arrow_traits<std::vector<byte_t>>::const_reference>;
 
    namespace detail
    {
        template <>
        struct get_data_type_from_array<sparrow::string_view_array>
        {
            [[nodiscard]] static constexpr sparrow::data_type get()
            {
                return sparrow::data_type::STRING_VIEW;
            }
        };
 
        template <>
        struct get_data_type_from_array<sparrow::binary_view_array>
        {
            [[nodiscard]] static constexpr sparrow::data_type get()
            {
                return sparrow::data_type::BINARY_VIEW;
            }
        };
    }
 
    template <std::ranges::sized_range T, class CR>
    struct array_inner_types<variable_size_binary_view_array_impl<T, CR>> : array_inner_types_base
    {
        using array_type = variable_size_binary_view_array_impl<T, CR>;
        using inner_value_type = T;
        using inner_reference = CR;
        using inner_const_reference = inner_reference;
 
        using value_iterator = functor_index_iterator<detail::layout_value_functor<array_type, inner_reference>>;
        using const_value_iterator = functor_index_iterator<
            detail::layout_value_functor<const array_type, inner_const_reference>>;
        using iterator_tag = std::random_access_iterator_tag;
    };
 
    template <class T>
    struct is_variable_size_binary_view_array_impl : std::false_type
    {
    };
 
    template <std::ranges::sized_range T, class CR>
    struct is_variable_size_binary_view_array_impl<variable_size_binary_view_array_impl<T, CR>> : std::true_type
    {
    };

    template <class T>
    constexpr bool is_variable_size_binary_view_array = is_variable_size_binary_view_array_impl<T>::value;

    template <std::ranges::sized_range T, class CR>
    class variable_size_binary_view_array_impl final
        : public mutable_array_bitmap_base<variable_size_binary_view_array_impl<T, CR>>
    {
    public:
 
        using self_type = variable_size_binary_view_array_impl<T, CR>;
        using base_type = mutable_array_bitmap_base<self_type>;
 
        using inner_types = array_inner_types<self_type>;
        using inner_value_type = typename inner_types::inner_value_type;
        using inner_reference = typename inner_types::inner_reference;
        using inner_const_reference = typename inner_types::inner_const_reference;
 
        using bitmap_type = typename base_type::bitmap_type;
        using bitmap_reference = typename base_type::bitmap_reference;
        using bitmap_const_reference = typename base_type::bitmap_const_reference;
        using bitmap_iterator = typename base_type::bitmap_iterator;
        using const_bitmap_iterator = typename base_type::const_bitmap_iterator;
        using bitmap_range = typename base_type::bitmap_range;
        using const_bitmap_range = typename base_type::const_bitmap_range;
 
        using value_type = nullable<inner_value_type>;
        using reference = nullable<inner_reference, bitmap_reference>;
        using const_reference = nullable<inner_const_reference, bitmap_const_reference>;
 
        using size_type = typename base_type::size_type;
        using difference_type = typename base_type::difference_type;
        using iterator_tag = typename base_type::iterator_tag;
 
        using value_iterator = typename base_type::value_iterator;
        using const_value_iterator = typename base_type::const_value_iterator;
 
        using iterator = typename base_type::iterator;
        using const_iterator = typename base_type::const_iterator;

        explicit variable_size_binary_view_array_impl(arrow_proxy);

        template <class... Args>
            requires(mpl::excludes_copy_and_move_ctor_v<variable_size_binary_view_array_impl<T, CR>, Args...>)
        explicit variable_size_binary_view_array_impl(Args&&... args)
            : variable_size_binary_view_array_impl(create_proxy(std::forward<Args>(args)...))
        {
        }
 
    private:

        struct buffers
        {
            buffer<uint8_t> length_buffer;        
            buffer<uint8_t> long_string_storage;  
            u8_buffer<int64_t> buffer_sizes;      
        };

        [[nodiscard]] static constexpr std::string_view get_arrow_format()
        {
            return std::is_same_v<T, arrow_traits<std::string>::value_type> ? std::string_view("vu")
                                                                            : std::string_view("vz");
        }

        template <input_metadata_container METADATA_RANGE>
        [[nodiscard]] static ArrowSchema create_arrow_schema(
            std::optional<std::string_view> name,
            std::optional<METADATA_RANGE> metadata,
            std::optional<std::unordered_set<sparrow::ArrowFlag>> flags
        )
        {
            constexpr repeat_view<bool> children_ownership(true, 0);
            return make_arrow_schema(
                get_arrow_format(),
                std::move(name),
                std::move(metadata),
                flags,
                nullptr,  // children
                children_ownership,
                nullptr,  // dictionary
                true
            );
        }

        template <std::ranges::input_range R>
            requires std::convertible_to<std::ranges::range_value_t<R>, T>
        static buffers create_buffers(R&& range);

        template <std::ranges::input_range R, validity_bitmap_input VB = validity_bitmap, input_metadata_container METADATA_RANGE>
            requires std::convertible_to<std::ranges::range_value_t<R>, T>
        [[nodiscard]] static arrow_proxy create_proxy(
            R&& range,
            VB&& bitmap_input = validity_bitmap{},
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        );

        template <std::ranges::input_range NULLABLE_RANGE, input_metadata_container METADATA_RANGE>
            requires std::convertible_to<std::ranges::range_value_t<NULLABLE_RANGE>, nullable<T>>
        [[nodiscard]] static arrow_proxy create_proxy(
            NULLABLE_RANGE&& nullable_range,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        );

        template <std::ranges::input_range R, input_metadata_container METADATA_RANGE>
            requires std::convertible_to<std::ranges::range_value_t<R>, T>
        [[nodiscard]] static arrow_proxy create_proxy(
            R&& range,
            bool = true,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        );

        template <std::ranges::input_range VALUE_BUFFERS_RANGE, validity_bitmap_input VB, input_metadata_container METADATA_RANGE>
            requires std::convertible_to<std::ranges::range_value_t<VALUE_BUFFERS_RANGE>, u8_buffer<uint8_t>>
        [[nodiscard]] static arrow_proxy create_proxy(
            size_t element_count,
            u8_buffer<uint8_t>&& buffer_view,
            VALUE_BUFFERS_RANGE&& value_buffers,
            VB&& validity_input,
            std::optional<std::string_view> name = std::nullopt,
            std::optional<METADATA_RANGE> metadata = std::nullopt
        );

        [[nodiscard]] constexpr inner_reference value(size_type i);

        [[nodiscard]] constexpr inner_const_reference value(size_type i) const;

        [[nodiscard]] constexpr value_iterator value_begin();

        [[nodiscard]] constexpr value_iterator value_end();

        [[nodiscard]] constexpr const_value_iterator value_cbegin() const;

        [[nodiscard]] constexpr const_value_iterator value_cend() const;
 
        static constexpr size_type LENGTH_BUFFER_INDEX = 1;            
        static constexpr std::size_t DATA_BUFFER_SIZE = 16;            
        static constexpr std::size_t SHORT_STRING_SIZE = 12;           
        static constexpr std::size_t PREFIX_SIZE = 4;                  
        static constexpr std::ptrdiff_t PREFIX_OFFSET = 4;             
        static constexpr std::ptrdiff_t SHORT_STRING_OFFSET = 4;       
        static constexpr std::ptrdiff_t BUFFER_INDEX_OFFSET = 8;       
        static constexpr std::ptrdiff_t BUFFER_OFFSET_OFFSET = 12;     
        static constexpr std::size_t FIRST_VAR_DATA_BUFFER_INDEX = 2;  
 
        friend base_type;
        friend base_type::base_type;
        friend base_type::base_type::base_type;
        friend class detail::layout_value_functor<self_type, inner_reference>;
        friend class detail::layout_value_functor<const self_type, inner_const_reference>;
    };
 
    template <std::ranges::sized_range T, class CR>
    variable_size_binary_view_array_impl<T, CR>::variable_size_binary_view_array_impl(arrow_proxy proxy)
        : base_type(std::move(proxy))
    {
    }
 
    template <std::ranges::sized_range T, class CR>
    template <std::ranges::input_range R>
        requires std::convertible_to<std::ranges::range_value_t<R>, T>
    auto variable_size_binary_view_array_impl<T, CR>::create_buffers(R&& range) -> buffers
    {
#ifdef __GNUC__
#    pragma GCC diagnostic push
#    pragma GCC diagnostic ignored "-Wcast-align"
#endif
 
        // Helper lambda to cast values to uint8_t
        auto to_uint8 = [](const auto& v)
        {
            return static_cast<std::uint8_t>(v);
        };
 
        const auto size = range_size(range);
        buffer<uint8_t> length_buffer(size * DATA_BUFFER_SIZE);
 
        std::size_t long_string_storage_size = 0;
        std::size_t i = 0;
        for (auto&& val : range)
        {
            auto val_casted = val | std::ranges::views::transform(to_uint8);
 
            const auto length = val.size();
            auto length_ptr = length_buffer.data() + (i * DATA_BUFFER_SIZE);
 
            // write length
            *reinterpret_cast<std::int32_t*>(length_ptr) = static_cast<std::int32_t>(length);
 
            if (length <= SHORT_STRING_SIZE)
            {
                // write data itself
                sparrow::ranges::copy(val_casted, length_ptr + SHORT_STRING_OFFSET);
                std::fill(
                    length_ptr + SHORT_STRING_OFFSET + length,
                    length_ptr + DATA_BUFFER_SIZE,
                    std::uint8_t(0)
                );
            }
            else
            {
                // write the prefix of the data
                auto prefix_sub_range = val_casted | std::ranges::views::take(PREFIX_SIZE);
                sparrow::ranges::copy(prefix_sub_range, length_ptr + PREFIX_OFFSET);
 
                // write the buffer index
                *reinterpret_cast<std::int32_t*>(length_ptr + BUFFER_INDEX_OFFSET) = 0;
 
                // write the buffer offset
                *reinterpret_cast<std::int32_t*>(
                    length_ptr + BUFFER_OFFSET_OFFSET
                ) = static_cast<std::int32_t>(long_string_storage_size);
 
                // count the size of the long string storage
                long_string_storage_size += length;
            }
            ++i;
        }
 
        // write the long string storage
        buffer<uint8_t> long_string_storage(long_string_storage_size);
        std::size_t long_string_storage_offset = 0;
        for (auto&& val : range)
        {
            const auto length = val.size();
            if (length > SHORT_STRING_SIZE)
            {
                auto val_casted = val | std::ranges::views::transform(to_uint8);
 
                sparrow::ranges::copy(val_casted, long_string_storage.data() + long_string_storage_offset);
                long_string_storage_offset += length;
            }
        }
 
        // For binary or utf-8 view arrays, an extra buffer is appended which stores
        // the lengths of each variadic data buffer as int64_t.
        // This buffer is necessary since these buffer lengths are not trivially
        // extractable from other data in an array of binary or utf-8 view type.
        u8_buffer<int64_t> buffer_sizes(
            static_cast<std::size_t>(1),
            static_cast<int64_t>(long_string_storage_size)
        );
 
        return {std::move(length_buffer), std::move(long_string_storage), std::move(buffer_sizes)};
 
#ifdef __GNUC__
#    pragma GCC diagnostic pop
#endif
    }
 
    template <std::ranges::sized_range T, class CR>
    template <std::ranges::input_range R, validity_bitmap_input VB, input_metadata_container METADATA_RANGE>
        requires std::convertible_to<std::ranges::range_value_t<R>, T>
    arrow_proxy variable_size_binary_view_array_impl<T, CR>::create_proxy(
        R&& range,
        VB&& validity_input,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    )
    {
        const auto size = range_size(range);
        validity_bitmap vbitmap = ensure_validity_bitmap(size, std::forward<VB>(validity_input));
        const auto null_count = vbitmap.null_count();
 
        static const std::optional<std::unordered_set<sparrow::ArrowFlag>> flags{{ArrowFlag::NULLABLE}};
 
        // create arrow schema
        ArrowSchema schema = create_arrow_schema(std::move(name), std::move(metadata), flags);
 
        // create buffers
        auto buffers_parts = create_buffers(std::forward<R>(range));
 
        std::vector<buffer<uint8_t>> buffers{
            std::move(vbitmap).extract_storage(),
            std::move(buffers_parts.length_buffer),
            std::move(buffers_parts.long_string_storage),
            std::move(buffers_parts.buffer_sizes).extract_storage()
        };
 
        constexpr repeat_view<bool> children_ownership(true, 0);
 
        // create arrow array
        ArrowArray arr = make_arrow_array(
            static_cast<std::int64_t>(size),  // length
            static_cast<int64_t>(null_count),
            0,  // offset
            std::move(buffers),
            nullptr,  // children
            children_ownership,
            nullptr,  // dictionary
            true
        );
 
        return arrow_proxy{std::move(arr), std::move(schema)};
    }
 
    template <std::ranges::sized_range T, class CR>
    template <std::ranges::input_range NULLABLE_RANGE, input_metadata_container METADATA_RANGE>
        requires std::convertible_to<std::ranges::range_value_t<NULLABLE_RANGE>, nullable<T>>
    [[nodiscard]] arrow_proxy variable_size_binary_view_array_impl<T, CR>::create_proxy(
        NULLABLE_RANGE&& nullable_range,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    )
    {
        auto values = nullable_range
                      | std::views::transform(
                          [](const auto& v)
                          {
                              return static_cast<T>(v.value());
                          }
                      );
 
        auto is_non_null = nullable_range
                           | std::views::transform(
                               [](const auto& v)
                               {
                                   return v.has_value();
                               }
                           );
 
        return create_proxy(
            std::forward<decltype(values)>(values),
            std::forward<decltype(is_non_null)>(is_non_null),
            name,
            metadata
        );
    }
 
    template <std::ranges::sized_range T, class CR>
    template <std::ranges::input_range R, input_metadata_container METADATA_RANGE>
        requires std::convertible_to<std::ranges::range_value_t<R>, T>
    [[nodiscard]] arrow_proxy variable_size_binary_view_array_impl<T, CR>::create_proxy(
        R&& range,
        bool nullable,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    )
    {
        if (nullable)
        {
            return create_proxy(std::forward<R>(range), validity_bitmap{}, std::move(name), std::move(metadata));
        }
 
        // create arrow schema
        ArrowSchema schema = create_arrow_schema(std::move(name), std::move(metadata), std::nullopt);
 
        // create buffers
        auto buffers_parts = create_buffers(std::forward<R>(range));
 
        std::vector<buffer<uint8_t>> buffers{
            buffer<uint8_t>{nullptr, 0},  // validity bitmap
            std::move(buffers_parts.length_buffer),
            std::move(buffers_parts.long_string_storage),
            std::move(buffers_parts.buffer_sizes).extract_storage()
        };
        const auto size = range_size(range);
 
        constexpr repeat_view<bool> children_ownership(true, 0);
 
        // create arrow array
        ArrowArray arr = make_arrow_array(
            static_cast<std::int64_t>(size),  // length
            static_cast<int64_t>(0),
            0,  // offset
            std::move(buffers),
            nullptr,  // children
            children_ownership,
            nullptr,  // dictionary
            true
        );
 
        return arrow_proxy{std::move(arr), std::move(schema)};
    }
 
    template <std::ranges::sized_range T, class CR>
    template <std::ranges::input_range VALUE_BUFFERS_RANGE, validity_bitmap_input VB, input_metadata_container METADATA_RANGE>
        requires std::convertible_to<std::ranges::range_value_t<VALUE_BUFFERS_RANGE>, u8_buffer<uint8_t>>
    arrow_proxy variable_size_binary_view_array_impl<T, CR>::create_proxy(
        size_t element_count,
        u8_buffer<uint8_t>&& buffer_view,
        VALUE_BUFFERS_RANGE&& value_buffers,
        VB&& validity_input,
        std::optional<std::string_view> name,
        std::optional<METADATA_RANGE> metadata
    )
    {
        const auto size = buffer_view.size() / DATA_BUFFER_SIZE;
        SPARROW_ASSERT_TRUE(size == element_count);
 
        static const std::optional<std::unordered_set<sparrow::ArrowFlag>> flags{{ArrowFlag::NULLABLE}};
 
        ArrowSchema schema = create_arrow_schema(std::move(name), std::move(metadata), flags);
 
        auto bitmap = ensure_validity_bitmap(size, std::forward<VB>(validity_input));
        std::vector<buffer<uint8_t>> buffers{std::move(bitmap).extract_storage(), std::move(buffer_view)};
        for (auto&& buf : value_buffers)
        {
            buffers.push_back(std::forward<decltype(buf)>(buf));
        }
 
        // Create buffer sizes for the variadic buffers
        u8_buffer<int64_t> buffer_sizes(value_buffers.size());
        for (std::size_t i = 0; i < value_buffers.size(); ++i)
        {
            buffer_sizes[i] = static_cast<int64_t>(value_buffers[i].size());
        }
        buffers.push_back(std::move(buffer_sizes).extract_storage());
 
        constexpr repeat_view<bool> children_ownership(true, 0);
 
        ArrowArray arr = make_arrow_array(
            static_cast<std::int64_t>(size),                 // length
            static_cast<std::int64_t>(bitmap.null_count()),  // null_count
            0,                                               // offset
            std::move(buffers),
            nullptr,  // children
            children_ownership,
            nullptr,  // dictionary
            true
        );
 
        return arrow_proxy{std::move(arr), std::move(schema)};
    }
 
    template <std::ranges::sized_range T, class CR>
    constexpr auto variable_size_binary_view_array_impl<T, CR>::value(size_type i) -> inner_reference
    {
        return static_cast<const self_type*>(this)->value(i);
    }
 
    template <std::ranges::sized_range T, class CR>
    constexpr auto variable_size_binary_view_array_impl<T, CR>::value(size_type i) const
        -> inner_const_reference
    {
#ifdef __GNUC__
#    pragma GCC diagnostic push
#    pragma GCC diagnostic ignored "-Wcast-align"
#endif
 
        SPARROW_ASSERT_TRUE(i < this->size());
        using char_or_byte = typename inner_const_reference::value_type;
 
        auto data_ptr = this->get_arrow_proxy().buffers()[LENGTH_BUFFER_INDEX].template data<uint8_t>()
                        + (i * DATA_BUFFER_SIZE);
        const auto length = static_cast<std::size_t>(*reinterpret_cast<const std::int32_t*>(data_ptr));
 
        if (length <= SHORT_STRING_SIZE)
        {
            constexpr std::ptrdiff_t data_offset = 4;
            const auto ptr = reinterpret_cast<const char_or_byte*>(data_ptr);
            const auto ret = inner_const_reference(ptr + data_offset, length);
            return ret;
        }
        else
        {
            const auto buffer_index = static_cast<std::size_t>(
                *reinterpret_cast<const std::int32_t*>(data_ptr + BUFFER_INDEX_OFFSET)
            );
            const auto buffer_offset = static_cast<std::size_t>(
                *reinterpret_cast<const std::int32_t*>(data_ptr + BUFFER_OFFSET_OFFSET)
            );
            const auto buffer = this->get_arrow_proxy()
                                    .buffers()[buffer_index + FIRST_VAR_DATA_BUFFER_INDEX]
                                    .template data<const char_or_byte>();
            return inner_const_reference(buffer + buffer_offset, length);
        }
 
#ifdef __GNUC__
#    pragma GCC diagnostic pop
#endif
    }
 
    template <std::ranges::sized_range T, class CR>
    constexpr auto variable_size_binary_view_array_impl<T, CR>::value_begin() -> value_iterator
    {
        return value_iterator(detail::layout_value_functor<self_type, inner_reference>(), 0);
    }
 
    template <std::ranges::sized_range T, class CR>
    constexpr auto variable_size_binary_view_array_impl<T, CR>::value_end() -> value_iterator
    {
        return value_iterator(detail::layout_value_functor<self_type, inner_reference>(this), this->size());
    }
 
    template <std::ranges::sized_range T, class CR>
    constexpr auto variable_size_binary_view_array_impl<T, CR>::value_cbegin() const -> const_value_iterator
    {
        return const_value_iterator(detail::layout_value_functor<const self_type, inner_const_reference>(this), 0);
    }
 
    template <std::ranges::sized_range T, class CR>
    constexpr auto variable_size_binary_view_array_impl<T, CR>::value_cend() const -> const_value_iterator
    {
        return const_value_iterator(
            detail::layout_value_functor<const self_type, inner_const_reference>(this),
            this->size()
        );
    }
}