#ifndef ARROW_ARRAY_H

#define ARROW_ARRAY_H

#include <cstdint>

#include <iosfwd>

#include <memory>

#include <string>

#include <type_traits>

#include <utility>

#include <vector>

#include "arrow/buffer.h"

#include "arrow/compare.h"

#include "arrow/result.h"

#include "arrow/type.h"

#include "arrow/type_traits.h"

#include "arrow/util/bit_util.h"

#include "arrow/util/checked_cast.h"

#include "arrow/util/macros.h"

#include "arrow/util/string_view.h" // IWYU pragma: export

#include "arrow/util/visibility.h"

namespace arrow {

class Array;

class ArrayVisitor;

constexpr int64_t kUnknownNullCount = -1;

class MemoryPool;

class Status;

struct ARROW_EXPORT ArrayData {

ArrayData() : length(0), null_count(0), offset(0) {}

ArrayData(const std::shared_ptr<DataType>& type, int64_t length,

int64_t null_count = kUnknownNullCount, int64_t offset = 0)

: type(type), length(length), null_count(null_count), offset(offset) {}

ArrayData(const std::shared_ptr<DataType>& type, int64_t length,

}

ArrayData(const std::shared_ptr<DataType>& type, int64_t length,

}

ArrayData(const std::shared_ptr<DataType>& type, int64_t length,

}

static std::shared_ptr<ArrayData> Make(const std::shared_ptr<DataType>& type,

int64_t length,

std::vector<std::shared_ptr<Buffer>>&& buffers,

int64_t null_count = kUnknownNullCount,

int64_t offset = 0);

static std::shared_ptr<ArrayData> Make(

const std::shared_ptr<DataType>& type, int64_t length,

const std::vector<std::shared_ptr<Buffer>>& buffers,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

static std::shared_ptr<ArrayData> Make(

const std::shared_ptr<DataType>& type, int64_t length,

const std::vector<std::shared_ptr<Buffer>>& buffers,

const std::vector<std::shared_ptr<ArrayData>>& child_data,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

static std::shared_ptr<ArrayData> Make(const std::shared_ptr<DataType>& type,

int64_t length,

int64_t null_count = kUnknownNullCount,

int64_t offset = 0);

// Move constructor

ArrayData(ArrayData&& other) noexcept

dictionary(std::move(other.dictionary)) {}

// Copy constructor

ArrayData(const ArrayData& other) noexcept

dictionary(other.dictionary) {}

// Move assignment

ArrayData& operator=(ArrayData&& other) = default;

// Copy assignment

ArrayData& operator=(const ArrayData& other) = default;

std::shared_ptr<ArrayData> Copy() const { return std::make_shared<ArrayData>(*this); }

// Access a buffer's data as a typed C pointer

template <typename T>

inline const T* GetValues(int i, int64_t absolute_offset) const {

if (buffers[i]) {

return reinterpret_cast<const T*>(buffers[i]->data()) + absolute_offset;

} else {

return NULLPTR;

}

}

template <typename T>

inline const T* GetValues(int i) const {

return GetValues<T>(i, offset);

}

// Access a buffer's data as a typed C pointer

template <typename T>

inline T* GetMutableValues(int i, int64_t absolute_offset) {

if (buffers[i]) {

return reinterpret_cast<T*>(buffers[i]->mutable_data()) + absolute_offset;

} else {

return NULLPTR;

}

}

template <typename T>

inline T* GetMutableValues(int i) {

return GetMutableValues<T>(i, offset);

}

// Construct a zero-copy slice of the data with the indicated offset and length

ArrayData Slice(int64_t offset, int64_t length) const;

/// \brief Return null count, or compute and set it if it's not known

int64_t GetNullCount() const;

std::shared_ptr<DataType> type;

int64_t length;

mutable int64_t null_count;

// The logical start point into the physical buffers (in values, not bytes).

// Note that, for child data, this must be *added* to the child data's own offset.

int64_t offset;

std::vector<std::shared_ptr<Buffer>> buffers;

std::vector<std::shared_ptr<ArrayData>> child_data;

// The dictionary for this Array, if any. Only used for dictionary

// type

std::shared_ptr<Array> dictionary;

};

std::shared_ptr<Array> MakeArray(const std::shared_ptr<ArrayData>& data);

Status MakeArrayOfNull(const std::shared_ptr<DataType>& type, int64_t length,

std::shared_ptr<Array>* out);

class ARROW_EXPORT Array {

public:

virtual ~Array() = default;

/// \brief Return true if value at index is null. Does not boundscheck

bool IsNull(int64_t i) const {

return null_bitmap_data_ != NULLPTR &&

!BitUtil::GetBit(null_bitmap_data_, i + data_->offset);

}

/// \brief Return true if value at index is valid (not null). Does not

/// boundscheck

bool IsValid(int64_t i) const {

return null_bitmap_data_ == NULLPTR ||

BitUtil::GetBit(null_bitmap_data_, i + data_->offset);

}

/// Size in the number of elements this array contains.

int64_t length() const { return data_->length; }

/// A relative position into another array's data, to enable zero-copy

/// slicing. This value defaults to zero

int64_t offset() const { return data_->offset; }

/// The number of null entries in the array. If the null count was not known

/// at time of construction (and set to a negative value), then the null

/// count will be computed and cached on the first invocation of this

/// function

int64_t null_count() const;

std::shared_ptr<DataType> type() const { return data_->type; }

Type::type type_id() const { return data_->type->id(); }

/// Buffer for the null bitmap.

///

/// Note that for `null_count == 0`, this can be null.

/// This buffer does not account for any slice offset

std::shared_ptr<Buffer> null_bitmap() const { return data_->buffers[0]; }

/// Raw pointer to the null bitmap.

///

/// Note that for `null_count == 0`, this can be null.

/// This buffer does not account for any slice offset

const uint8_t* null_bitmap_data() const { return null_bitmap_data_; }

/// Equality comparison with another array

bool Equals(const Array& arr, const EqualOptions& = EqualOptions::Defaults()) const;

bool Equals(const std::shared_ptr<Array>& arr,

const EqualOptions& = EqualOptions::Defaults()) const;

/// Approximate equality comparison with another array

///

/// epsilon is only used if this is FloatArray or DoubleArray

bool ApproxEquals(const std::shared_ptr<Array>& arr,

const EqualOptions& = EqualOptions::Defaults()) const;

bool ApproxEquals(const Array& arr,

const EqualOptions& = EqualOptions::Defaults()) const;

/// Compare if the range of slots specified are equal for the given array and

/// this array. end_idx exclusive. This methods does not bounds check.

bool RangeEquals(int64_t start_idx, int64_t end_idx, int64_t other_start_idx,

const Array& other) const;

bool RangeEquals(int64_t start_idx, int64_t end_idx, int64_t other_start_idx,

const std::shared_ptr<Array>& other) const;

bool RangeEquals(const Array& other, int64_t start_idx, int64_t end_idx,

int64_t other_start_idx) const;

bool RangeEquals(const std::shared_ptr<Array>& other, int64_t start_idx,

int64_t end_idx, int64_t other_start_idx) const;

Status Accept(ArrayVisitor* visitor) const;

/// Construct a zero-copy view of this array with the given type.

///

/// This method checks if the types are layout-compatible.

/// Nested types are traversed in depth-first order. Data buffers must have

/// the same item sizes, even though the logical types may be different.

/// An error is returned if the types are not layout-compatible.

Status View(const std::shared_ptr<DataType>& type, std::shared_ptr<Array>* out) const;

/// Construct a zero-copy slice of the array with the indicated offset and

/// length

///

/// \param[in] offset the position of the first element in the constructed

/// slice

/// \param[in] length the length of the slice. If there are not enough

/// elements in the array, the length will be adjusted accordingly

///

/// \return a new object wrapped in std::shared_ptr<Array>

std::shared_ptr<Array> Slice(int64_t offset, int64_t length) const;

/// Slice from offset until end of the array

std::shared_ptr<Array> Slice(int64_t offset) const;

std::shared_ptr<ArrayData> data() const { return data_; }

int num_fields() const { return static_cast<int>(data_->child_data.size()); }

/// \return PrettyPrint representation of array suitable for debugging

std::string ToString() const;

/// \brief Perform any validation checks to determine obvious inconsistencies

/// within the array's internal data.

///

/// This can be an expensive check, potentially O(length).

///

/// \return Status

Status Validate() const;

protected:

Array() : null_bitmap_data_(NULLPTR) {}

std::shared_ptr<ArrayData> data_;

const uint8_t* null_bitmap_data_;

/// Protected method for constructors

inline void SetData(const std::shared_ptr<ArrayData>& data) {

}

private:

ARROW_DISALLOW_COPY_AND_ASSIGN(Array);

};

using ArrayVector = std::vector<std::shared_ptr<Array>>;

namespace internal {

std::vector<ArrayVector> RechunkArraysConsistently(const std::vector<ArrayVector>&);

} // namespace internal

static inline std::ostream& operator<<(std::ostream& os, const Array& x) {

os << x.ToString();

return os;

}

class ARROW_EXPORT FlatArray : public Array {

protected:

using Array::Array;

};

class ARROW_EXPORT NullArray : public FlatArray {

public:

using TypeClass = NullType;

explicit NullArray(const std::shared_ptr<ArrayData>& data) { SetData(data); }

explicit NullArray(int64_t length);

private:

inline void SetData(const std::shared_ptr<ArrayData>& data) {

null_bitmap_data_ = NULLPTR;

data->null_count = data->length;

data_ = data;

}

};

class ARROW_EXPORT PrimitiveArray : public FlatArray {

public:

PrimitiveArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::shared_ptr<Buffer>& data,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

/// Does not account for any slice offset

std::shared_ptr<Buffer> values() const { return data_->buffers[1]; }

protected:

PrimitiveArray() : raw_values_(NULLPTR) {}

inline void SetData(const std::shared_ptr<ArrayData>& data) {

auto values = data->buffers[1];

this->Array::SetData(data);

raw_values_ = values == NULLPTR ? NULLPTR : values->data();

}

explicit inline PrimitiveArray(const std::shared_ptr<ArrayData>& data) {

SetData(data);

}

const uint8_t* raw_values_;

};

template <typename TYPE>

class NumericArray : public PrimitiveArray {

};

class ARROW_EXPORT BooleanArray : public PrimitiveArray {

public:

using TypeClass = BooleanType;

explicit BooleanArray(const std::shared_ptr<ArrayData>& data);

BooleanArray(int64_t length, const std::shared_ptr<Buffer>& data,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

bool Value(int64_t i) const {

return BitUtil::GetBit(reinterpret_cast<const uint8_t*>(raw_values_),

i + data_->offset);

}

bool GetView(int64_t i) const { return Value(i); }

protected:

using PrimitiveArray::PrimitiveArray;

};

template <typename TYPE>

class BaseListArray : public Array {

};

class ARROW_EXPORT ListArray : public BaseListArray<ListType> {

public:

explicit ListArray(const std::shared_ptr<ArrayData>& data);

ListArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::shared_ptr<Buffer>& value_offsets,

const std::shared_ptr<Array>& values,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

/// \brief Construct ListArray from array of offsets and child value array

///

/// This function does the bare minimum of validation of the offsets and

/// input types, and will allocate a new offsets array if necessary (i.e. if

/// the offsets contain any nulls). If the offsets do not have nulls, they

/// are assumed to be well-formed

///

/// \param[in] offsets Array containing n + 1 offsets encoding length and

/// size. Must be of int32 type

/// \param[in] values Array containing list values

/// \param[in] pool MemoryPool in case new offsets array needs to be

/// allocated because of null values

/// \param[out] out Will have length equal to offsets.length() - 1

static Status FromArrays(const Array& offsets, const Array& values, MemoryPool* pool,

std::shared_ptr<Array>* out);

protected:

// This constructor defers SetData to a derived array class

ListArray() = default;

void SetData(const std::shared_ptr<ArrayData>& data);

};

class ARROW_EXPORT LargeListArray : public BaseListArray<LargeListType> {

public:

explicit LargeListArray(const std::shared_ptr<ArrayData>& data);

LargeListArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::shared_ptr<Buffer>& value_offsets,

const std::shared_ptr<Array>& values,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

/// \brief Construct LargeListArray from array of offsets and child value array

///

/// This function does the bare minimum of validation of the offsets and

/// input types, and will allocate a new offsets array if necessary (i.e. if

/// the offsets contain any nulls). If the offsets do not have nulls, they

/// are assumed to be well-formed

///

/// \param[in] offsets Array containing n + 1 offsets encoding length and

/// size. Must be of int64 type

/// \param[in] values Array containing list values

/// \param[in] pool MemoryPool in case new offsets array needs to be

/// allocated because of null values

/// \param[out] out Will have length equal to offsets.length() - 1

static Status FromArrays(const Array& offsets, const Array& values, MemoryPool* pool,

std::shared_ptr<Array>* out);

protected:

void SetData(const std::shared_ptr<ArrayData>& data);

};

class ARROW_EXPORT MapArray : public ListArray {

public:

using TypeClass = MapType;

explicit MapArray(const std::shared_ptr<ArrayData>& data);

MapArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::shared_ptr<Buffer>& value_offsets,

const std::shared_ptr<Array>& keys, const std::shared_ptr<Array>& items,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

MapArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::shared_ptr<Buffer>& value_offsets,

const std::shared_ptr<Array>& values,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

const MapType* map_type() const { return map_type_; }

/// \brief Return array object containing all map keys

std::shared_ptr<Array> keys() const { return keys_; }

/// \brief Return array object containing all mapped items

std::shared_ptr<Array> items() const { return items_; }

protected:

void SetData(const std::shared_ptr<ArrayData>& data);

private:

const MapType* map_type_;

std::shared_ptr<Array> keys_, items_;

};

class ARROW_EXPORT FixedSizeListArray : public Array {

public:

using TypeClass = FixedSizeListType;

using offset_type = TypeClass::offset_type;

explicit FixedSizeListArray(const std::shared_ptr<ArrayData>& data);

FixedSizeListArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::shared_ptr<Array>& values,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

const FixedSizeListType* list_type() const;

/// \brief Return array object containing the list's values

std::shared_ptr<Array> values() const;

std::shared_ptr<DataType> value_type() const;

// The following functions will not perform boundschecking

int32_t value_offset(int64_t i) const {

i += data_->offset;

return static_cast<int32_t>(list_size_ * i);

}

int32_t value_length(int64_t i = 0) const { return list_size_; }

std::shared_ptr<Array> value_slice(int64_t i) const {

return values_->Slice(value_offset(i), value_length(i));

}

protected:

void SetData(const std::shared_ptr<ArrayData>& data);

int32_t list_size_;

private:

std::shared_ptr<Array> values_;

};

template <typename TYPE>

class BaseBinaryArray : public FlatArray {

};

class ARROW_EXPORT BinaryArray : public BaseBinaryArray<BinaryType> {

public:

explicit BinaryArray(const std::shared_ptr<ArrayData>& data);

BinaryArray(int64_t length, const std::shared_ptr<Buffer>& value_offsets,

BinaryArray() : BaseBinaryArray() {}

};

class ARROW_EXPORT StringArray : public BinaryArray {

public:

using TypeClass = StringType;

explicit StringArray(const std::shared_ptr<ArrayData>& data);

StringArray(int64_t length, const std::shared_ptr<Buffer>& value_offsets,

const std::shared_ptr<Buffer>& data,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

};

class ARROW_EXPORT LargeBinaryArray : public BaseBinaryArray<LargeBinaryType> {

public:

explicit LargeBinaryArray(const std::shared_ptr<ArrayData>& data);

LargeBinaryArray(int64_t length, const std::shared_ptr<Buffer>& value_offsets,

LargeBinaryArray() : BaseBinaryArray() {}

};

class ARROW_EXPORT LargeStringArray : public LargeBinaryArray {

public:

using TypeClass = LargeStringType;

explicit LargeStringArray(const std::shared_ptr<ArrayData>& data);

LargeStringArray(int64_t length, const std::shared_ptr<Buffer>& value_offsets,

const std::shared_ptr<Buffer>& data,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

};

class ARROW_EXPORT FixedSizeBinaryArray : public PrimitiveArray {

public:

using TypeClass = FixedSizeBinaryType;

explicit FixedSizeBinaryArray(const std::shared_ptr<ArrayData>& data);

FixedSizeBinaryArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::shared_ptr<Buffer>& data,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

const uint8_t* GetValue(int64_t i) const;

const uint8_t* Value(int64_t i) const { return GetValue(i); }

util::string_view GetView(int64_t i) const {

return util::string_view(reinterpret_cast<const char*>(GetValue(i)), byte_width());

}

std::string GetString(int64_t i) const { return std::string(GetView(i)); }

int32_t byte_width() const { return byte_width_; }

const uint8_t* raw_values() const { return raw_values_ + data_->offset * byte_width_; }

protected:

inline void SetData(const std::shared_ptr<ArrayData>& data) {

this->PrimitiveArray::SetData(data);

byte_width_ =

internal::checked_cast<const FixedSizeBinaryType&>(*type()).byte_width();

}

int32_t byte_width_;

};

class ARROW_EXPORT DayTimeIntervalArray : public PrimitiveArray {

public:

using TypeClass = DayTimeIntervalType;

explicit DayTimeIntervalArray(const std::shared_ptr<ArrayData>& data);

DayTimeIntervalArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::shared_ptr<Buffer>& data,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

TypeClass::DayMilliseconds GetValue(int64_t i) const;

TypeClass::DayMilliseconds Value(int64_t i) const { return GetValue(i); }

// For compatibility with Take kernel.

TypeClass::DayMilliseconds GetView(int64_t i) const { return GetValue(i); }

int32_t byte_width() const { return sizeof(TypeClass::DayMilliseconds); }

const uint8_t* raw_values() const { return raw_values_ + data_->offset * byte_width(); }

protected:

inline void SetData(const std::shared_ptr<ArrayData>& data) {

this->PrimitiveArray::SetData(data);

}

};

class ARROW_EXPORT Decimal128Array : public FixedSizeBinaryArray {

public:

using TypeClass = Decimal128Type;

using FixedSizeBinaryArray::FixedSizeBinaryArray;

/// \brief Construct Decimal128Array from ArrayData instance

explicit Decimal128Array(const std::shared_ptr<ArrayData>& data);

std::string FormatValue(int64_t i) const;

};

using DecimalArray = Decimal128Array;

class ARROW_EXPORT StructArray : public Array {

public:

using TypeClass = StructType;

explicit StructArray(const std::shared_ptr<ArrayData>& data);

StructArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::vector<std::shared_ptr<Array>>& children,

std::shared_ptr<Buffer> null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

/// \brief Return a StructArray from child arrays and field names.

///

/// The length and data type are automatically inferred from the arguments.

/// There should be at least one child array.

static Result<std::shared_ptr<StructArray>> Make(

const std::vector<std::shared_ptr<Array>>& children,

const std::vector<std::string>& field_names,

std::shared_ptr<Buffer> null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

/// \brief Return a StructArray from child arrays and fields.

///

/// The length is automatically inferred from the arguments.

/// There should be at least one child array. This method does not

/// check that field types and child array types are consistent.

static Result<std::shared_ptr<StructArray>> Make(

const std::vector<std::shared_ptr<Array>>& children,

const std::vector<std::shared_ptr<Field>>& fields,

std::shared_ptr<Buffer> null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

const StructType* struct_type() const;

// Return a shared pointer in case the requestor desires to share ownership

// with this array. The returned array has its offset, length and null

// count adjusted.

std::shared_ptr<Array> field(int pos) const;

/// Returns null if name not found

std::shared_ptr<Array> GetFieldByName(const std::string& name) const;

/// \brief Flatten this array as a vector of arrays, one for each field

///

/// \param[in] pool The pool to allocate null bitmaps from, if necessary

/// \param[out] out The resulting vector of arrays

Status Flatten(MemoryPool* pool, ArrayVector* out) const;

private:

// For caching boxed child data

// XXX This is not handled in a thread-safe manner.

mutable std::vector<std::shared_ptr<Array>> boxed_fields_;

};

class ARROW_EXPORT UnionArray : public Array {

public:

using TypeClass = UnionType;

using type_id_t = uint8_t;

explicit UnionArray(const std::shared_ptr<ArrayData>& data);

UnionArray(const std::shared_ptr<DataType>& type, int64_t length,

const std::vector<std::shared_ptr<Array>>& children,

const std::shared_ptr<Buffer>& type_ids,

const std::shared_ptr<Buffer>& value_offsets = NULLPTR,

const std::shared_ptr<Buffer>& null_bitmap = NULLPTR,

int64_t null_count = kUnknownNullCount, int64_t offset = 0);

/// \brief Construct Dense UnionArray from types_ids, value_offsets and children

///

/// This function does the bare minimum of validation of the offsets and

/// input types. The value_offsets are assumed to be well-formed.

///

/// \param[in] type_ids An array of 8-bit signed integers, enumerated from

/// 0 corresponding to each type.

/// \param[in] value_offsets An array of signed int32 values indicating the

/// relative offset into the respective child array for the type in a given slot.

/// The respective offsets for each child value array must be in order / increasing.

/// \param[in] children Vector of children Arrays containing the data for each type.

/// \param[in] field_names Vector of strings containing the name of each field.

/// \param[in] type_codes Vector of type codes.

/// \param[out] out Will have length equal to value_offsets.length()

static Status MakeDense(const Array& type_ids, const Array& value_offsets,

const std::vector<std::shared_ptr<Array>>& children,

const std::vector<std::string>& field_names,

const std::vector<uint8_t>& type_codes,

std::shared_ptr<Array>* out);