venv/lib/python3.11/site-packages/pyarrow/include/arrow/util/converter.h

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you 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.

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "arrow/array.h"
#include "arrow/chunked_array.h"
#include "arrow/status.h"
#include "arrow/type.h"
#include "arrow/type_traits.h"
#include "arrow/util/checked_cast.h"
#include "arrow/visit_type_inline.h"

namespace arrow {
namespace internal {

template <typename BaseConverter, template <typename...> class ConverterTrait>
static Result<std::unique_ptr<BaseConverter>> MakeConverter(
    std::shared_ptr<DataType> type, typename BaseConverter::OptionsType options,
    MemoryPool* pool);

template <typename Input, typename Options>
class Converter {
 public:
  using Self = Converter<Input, Options>;
  using InputType = Input;
  using OptionsType = Options;

  virtual ~Converter() = default;

  Status Construct(std::shared_ptr<DataType> type, OptionsType options,
                   MemoryPool* pool) {
    type_ = std::move(type);
    options_ = std::move(options);
    return Init(pool);
  }

  virtual Status Append(InputType value) { return Status::NotImplemented("Append"); }

  virtual Status Extend(InputType values, int64_t size, int64_t offset = 0) {
    return Status::NotImplemented("Extend");
  }

  virtual Status ExtendMasked(InputType values, InputType mask, int64_t size,
                              int64_t offset = 0) {
    return Status::NotImplemented("ExtendMasked");
  }

  const std::shared_ptr<ArrayBuilder>& builder() const { return builder_; }

  const std::shared_ptr<DataType>& type() const { return type_; }

  OptionsType options() const { return options_; }

  bool may_overflow() const { return may_overflow_; }

  bool rewind_on_overflow() const { return rewind_on_overflow_; }

  virtual Status Reserve(int64_t additional_capacity) {
    return builder_->Reserve(additional_capacity);
  }

  Status AppendNull() { return builder_->AppendNull(); }

  virtual Result<std::shared_ptr<Array>> ToArray() { return builder_->Finish(); }

  virtual Result<std::shared_ptr<Array>> ToArray(int64_t length) {
    ARROW_ASSIGN_OR_RAISE(auto arr, this->ToArray());
    return arr->Slice(0, length);
  }

  virtual Result<std::shared_ptr<ChunkedArray>> ToChunkedArray() {
    ARROW_ASSIGN_OR_RAISE(auto array, ToArray());
    std::vector<std::shared_ptr<Array>> chunks = {std::move(array)};
    return std::make_shared<ChunkedArray>(chunks);
  }

 protected:
  virtual Status Init(MemoryPool* pool) { return Status::OK(); }

  std::shared_ptr<DataType> type_;
  std::shared_ptr<ArrayBuilder> builder_;
  OptionsType options_;
  bool may_overflow_ = false;
  bool rewind_on_overflow_ = false;
};

template <typename ArrowType, typename BaseConverter>
class PrimitiveConverter : public BaseConverter {
 public:
  using BuilderType = typename TypeTraits<ArrowType>::BuilderType;

 protected:
  Status Init(MemoryPool* pool) override {
    this->builder_ = std::make_shared<BuilderType>(this->type_, pool);
    // Narrow variable-sized binary types may overflow
    this->may_overflow_ = is_binary_like(this->type_->id());
    primitive_type_ = checked_cast<const ArrowType*>(this->type_.get());
    primitive_builder_ = checked_cast<BuilderType*>(this->builder_.get());
    return Status::OK();
  }

  const ArrowType* primitive_type_;
  BuilderType* primitive_builder_;
};

template <typename ArrowType, typename BaseConverter,
          template <typename...> class ConverterTrait>
class ListConverter : public BaseConverter {
 public:
  using BuilderType = typename TypeTraits<ArrowType>::BuilderType;
  using ConverterType = typename ConverterTrait<ArrowType>::type;

 protected:
  Status Init(MemoryPool* pool) override {
    list_type_ = checked_cast<const ArrowType*>(this->type_.get());
    ARROW_ASSIGN_OR_RAISE(value_converter_,
                          (MakeConverter<BaseConverter, ConverterTrait>(
                              list_type_->value_type(), this->options_, pool)));
    this->builder_ =
        std::make_shared<BuilderType>(pool, value_converter_->builder(), this->type_);
    list_builder_ = checked_cast<BuilderType*>(this->builder_.get());
    // Narrow list types may overflow
    this->may_overflow_ = this->rewind_on_overflow_ =
        sizeof(typename ArrowType::offset_type) < sizeof(int64_t);
    return Status::OK();
  }

  const ArrowType* list_type_;
  BuilderType* list_builder_;
  std::unique_ptr<BaseConverter> value_converter_;
};

template <typename BaseConverter, template <typename...> class ConverterTrait>
class StructConverter : public BaseConverter {
 public:
  using ConverterType = typename ConverterTrait<StructType>::type;

  Status Reserve(int64_t additional_capacity) override {
    ARROW_RETURN_NOT_OK(this->builder_->Reserve(additional_capacity));
    for (const auto& child : children_) {
      ARROW_RETURN_NOT_OK(child->Reserve(additional_capacity));
    }
    return Status::OK();
  }

 protected:
  Status Init(MemoryPool* pool) override {
    std::unique_ptr<BaseConverter> child_converter;
    std::vector<std::shared_ptr<ArrayBuilder>> child_builders;

    struct_type_ = checked_cast<const StructType*>(this->type_.get());
    for (const auto& field : struct_type_->fields()) {
      ARROW_ASSIGN_OR_RAISE(child_converter,
                            (MakeConverter<BaseConverter, ConverterTrait>(
                                field->type(), this->options_, pool)));
      this->may_overflow_ |= child_converter->may_overflow();
      this->rewind_on_overflow_ = this->may_overflow_;
      child_builders.push_back(child_converter->builder());
      children_.push_back(std::move(child_converter));
    }

    this->builder_ =
        std::make_shared<StructBuilder>(this->type_, pool, std::move(child_builders));
    struct_builder_ = checked_cast<StructBuilder*>(this->builder_.get());

    return Status::OK();
  }

  const StructType* struct_type_;
  StructBuilder* struct_builder_;
  std::vector<std::unique_ptr<BaseConverter>> children_;
};

template <typename ValueType, typename BaseConverter>
class DictionaryConverter : public BaseConverter {
 public:
  using BuilderType = DictionaryBuilder<ValueType>;

 protected:
  Status Init(MemoryPool* pool) override {
    std::unique_ptr<ArrayBuilder> builder;
    ARROW_RETURN_NOT_OK(MakeDictionaryBuilder(pool, this->type_, NULLPTR, &builder));
    this->builder_ = std::move(builder);
    this->may_overflow_ = false;
    dict_type_ = checked_cast<const DictionaryType*>(this->type_.get());
    value_type_ = checked_cast<const ValueType*>(dict_type_->value_type().get());
    value_builder_ = checked_cast<BuilderType*>(this->builder_.get());
    return Status::OK();
  }

  const DictionaryType* dict_type_;
  const ValueType* value_type_;
  BuilderType* value_builder_;
};

template <typename BaseConverter, template <typename...> class ConverterTrait>
struct MakeConverterImpl {
  template <typename T, typename ConverterType = typename ConverterTrait<T>::type>
  Status Visit(const T&) {
    out.reset(new ConverterType());
    return out->Construct(std::move(type), std::move(options), pool);
  }

  Status Visit(const DictionaryType& t) {
    switch (t.value_type()->id()) {
#define DICTIONARY_CASE(TYPE)                                                       \
  case TYPE::type_id:                                                               \
    out = std::make_unique<                                                         \
        typename ConverterTrait<DictionaryType>::template dictionary_type<TYPE>>(); \
    break;
      DICTIONARY_CASE(BooleanType);
      DICTIONARY_CASE(Int8Type);
      DICTIONARY_CASE(Int16Type);
      DICTIONARY_CASE(Int32Type);
      DICTIONARY_CASE(Int64Type);
      DICTIONARY_CASE(UInt8Type);
      DICTIONARY_CASE(UInt16Type);
      DICTIONARY_CASE(UInt32Type);
      DICTIONARY_CASE(UInt64Type);
      DICTIONARY_CASE(FloatType);
      DICTIONARY_CASE(DoubleType);
      DICTIONARY_CASE(BinaryType);
      DICTIONARY_CASE(StringType);
      DICTIONARY_CASE(FixedSizeBinaryType);
#undef DICTIONARY_CASE
      default:
        return Status::NotImplemented("DictionaryArray converter for type ", t.ToString(),
                                      " not implemented");
    }
    return out->Construct(std::move(type), std::move(options), pool);
  }

  Status Visit(const DataType& t) { return Status::NotImplemented(t.name()); }

  std::shared_ptr<DataType> type;
  typename BaseConverter::OptionsType options;
  MemoryPool* pool;
  std::unique_ptr<BaseConverter> out;
};

template <typename BaseConverter, template <typename...> class ConverterTrait>
static Result<std::unique_ptr<BaseConverter>> MakeConverter(
    std::shared_ptr<DataType> type, typename BaseConverter::OptionsType options,
    MemoryPool* pool) {
  MakeConverterImpl<BaseConverter, ConverterTrait> visitor{
      std::move(type), std::move(options), pool, NULLPTR};
  ARROW_RETURN_NOT_OK(VisitTypeInline(*visitor.type, &visitor));
  return std::move(visitor.out);
}

template <typename Converter>
class Chunker {
 public:
  using InputType = typename Converter::InputType;

  explicit Chunker(std::unique_ptr<Converter> converter)
      : converter_(std::move(converter)) {}

  Status Reserve(int64_t additional_capacity) {
    ARROW_RETURN_NOT_OK(converter_->Reserve(additional_capacity));
    reserved_ += additional_capacity;
    return Status::OK();
  }

  Status AppendNull() {
    auto status = converter_->AppendNull();
    if (ARROW_PREDICT_FALSE(status.IsCapacityError())) {
      if (converter_->builder()->length() == 0) {
        // Builder length == 0 means the individual element is too large to append.
        // In this case, no need to try again.
        return status;
      }
      ARROW_RETURN_NOT_OK(FinishChunk());
      return converter_->AppendNull();
    }
    ++length_;
    return status;
  }

  Status Append(InputType value) {
    auto status = converter_->Append(value);
    if (ARROW_PREDICT_FALSE(status.IsCapacityError())) {
      if (converter_->builder()->length() == 0) {
        return status;
      }
      ARROW_RETURN_NOT_OK(FinishChunk());
      return Append(value);
    }
    ++length_;
    return status;
  }

  Status Extend(InputType values, int64_t size, int64_t offset = 0) {
    while (offset < size) {
      auto length_before = converter_->builder()->length();
      auto status = converter_->Extend(values, size, offset);
      auto length_after = converter_->builder()->length();
      auto num_converted = length_after - length_before;

      offset += num_converted;
      length_ += num_converted;

      if (status.IsCapacityError()) {
        if (converter_->builder()->length() == 0) {
          // Builder length == 0 means the individual element is too large to append.
          // In this case, no need to try again.
          return status;
        } else if (converter_->rewind_on_overflow()) {
          // The list-like and binary-like conversion paths may raise  a capacity error,
          // we need to handle them differently. While the binary-like converters check
          // the capacity before append/extend the list-like converters just check after
          // append/extend. Thus depending on the implementation semantics we may need
          // to rewind (slice) the output chunk by one.
          length_ -= 1;
          offset -= 1;
        }
        ARROW_RETURN_NOT_OK(FinishChunk());
      } else if (!status.ok()) {
        return status;
      }
    }
    return Status::OK();
  }

  Status ExtendMasked(InputType values, InputType mask, int64_t size,
                      int64_t offset = 0) {
    while (offset < size) {
      auto length_before = converter_->builder()->length();
      auto status = converter_->ExtendMasked(values, mask, size, offset);
      auto length_after = converter_->builder()->length();
      auto num_converted = length_after - length_before;

      offset += num_converted;
      length_ += num_converted;

      if (status.IsCapacityError()) {
        if (converter_->builder()->length() == 0) {
          // Builder length == 0 means the individual element is too large to append.
          // In this case, no need to try again.
          return status;
        } else if (converter_->rewind_on_overflow()) {
          // The list-like and binary-like conversion paths may raise  a capacity error,
          // we need to handle them differently. While the binary-like converters check
          // the capacity before append/extend the list-like converters just check after
          // append/extend. Thus depending on the implementation semantics we may need
          // to rewind (slice) the output chunk by one.
          length_ -= 1;
          offset -= 1;
        }
        ARROW_RETURN_NOT_OK(FinishChunk());
      } else if (!status.ok()) {
        return status;
      }
    }
    return Status::OK();
  }

  Status FinishChunk() {
    ARROW_ASSIGN_OR_RAISE(auto chunk, converter_->ToArray(length_));
    chunks_.push_back(chunk);
    // Reserve space for the remaining items.
    // Besides being an optimization, it is also required if the converter's
    // implementation relies on unsafe builder methods in converter->Append().
    auto remaining = reserved_ - length_;
    Reset();
    return Reserve(remaining);
  }

  Result<std::shared_ptr<ChunkedArray>> ToChunkedArray() {
    ARROW_RETURN_NOT_OK(FinishChunk());
    return std::make_shared<ChunkedArray>(chunks_);
  }

 protected:
  void Reset() {
    converter_->builder()->Reset();
    length_ = 0;
    reserved_ = 0;
  }

  int64_t length_ = 0;
  int64_t reserved_ = 0;
  std::unique_ptr<Converter> converter_;
  std::vector<std::shared_ptr<Array>> chunks_;
};

template <typename T>
static Result<std::unique_ptr<Chunker<T>>> MakeChunker(std::unique_ptr<T> converter) {
  return std::make_unique<Chunker<T>>(std::move(converter));
}

}  // namespace internal
}  // namespace arrow