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

// Functions for pandas conversion via NumPy

#include "arrow/python/arrow_to_pandas.h"

#include "arrow/python/numpy_interop.h" // IWYU pragma: expand

#include <cmath>

#include <cstdint>

#include <iostream>

#include <memory>

#include <mutex>

#include <string>

#include <string_view>

#include <unordered_map>

#include <utility>

#include <vector>

#include "arrow/array.h"

#include "arrow/buffer.h"

#include "arrow/datum.h"

#include "arrow/status.h"

#include "arrow/table.h"

#include "arrow/type.h"

#include "arrow/type_traits.h"

#include "arrow/util/checked_cast.h"

#include "arrow/util/hashing.h"

#include "arrow/util/int_util.h"

#include "arrow/util/logging.h"

#include "arrow/util/macros.h"

#include "arrow/util/parallel.h"

#include "arrow/visit_type_inline.h"

#include "arrow/compute/api.h"

#include "arrow/python/arrow_to_python_internal.h"

#include "arrow/python/common.h"

#include "arrow/python/datetime.h"

#include "arrow/python/decimal.h"

#include "arrow/python/helpers.h"

#include "arrow/python/numpy_convert.h"

#include "arrow/python/numpy_internal.h"

#include "arrow/python/pyarrow.h"

#include "arrow/python/python_to_arrow.h"

#include "arrow/python/type_traits.h"

namespace arrow {

class MemoryPool;

using internal::checked_cast;

using internal::CheckIndexBounds;

using internal::OptionalParallelFor;

namespace py {

namespace {

// Fix options for conversion of an inner (child) array.

PandasOptions MakeInnerOptions(PandasOptions options) {

// Make sure conversion of inner dictionary arrays always returns an array,

// not a dict {'indices': array, 'dictionary': array, 'ordered': bool}

options.decode_dictionaries = true;

options.categorical_columns.reset();

options.strings_to_categorical = false;

// In ARROW-7723, we found as a result of ARROW-3789 that second

// through microsecond resolution tz-aware timestamps were being promoted to

// use the DATETIME_NANO_TZ conversion path, yielding a datetime64[ns] NumPy

// array in this function. PyArray_GETITEM returns datetime.datetime for

// units second through microsecond but PyLong for nanosecond (because

// datetime.datetime does not support nanoseconds).

// We force the object conversion to preserve the value of the timezone.

// Nanoseconds are returned as integers.

options.coerce_temporal_nanoseconds = false;

return options;

}

// ----------------------------------------------------------------------

// PyCapsule code for setting ndarray base to reference C++ object

struct ArrayCapsule {

std::shared_ptr<Array> array;

};

struct BufferCapsule {

std::shared_ptr<Buffer> buffer;

};

void ArrayCapsule_Destructor(PyObject* capsule) {

delete reinterpret_cast<ArrayCapsule*>(PyCapsule_GetPointer(capsule, "arrow::Array"));

}

void BufferCapsule_Destructor(PyObject* capsule) {

delete reinterpret_cast<BufferCapsule*>(PyCapsule_GetPointer(capsule, "arrow::Buffer"));

}

// ----------------------------------------------------------------------

// pandas 0.x DataFrame conversion internals

using internal::arrow_traits;

using internal::npy_traits;

bool IsUuidExtension(const DataType& type) {

if (type.id() != Type::EXTENSION) {

return false;

}

const auto& extension_type = checked_cast<const ExtensionType&>(type);

const auto& storage_type = *extension_type.storage_type();

return extension_type.extension_name() == "arrow.uuid" &&

storage_type.id() == Type::FIXED_SIZE_BINARY &&

checked_cast<const FixedSizeBinaryType&>(storage_type).byte_width() == 16;

}

template <typename T>

struct WrapBytes {};

template <>

struct WrapBytes<StringType> {

static inline PyObject* Wrap(const char* data, int64_t length) {

return PyUnicode_FromStringAndSize(data, length);

}

};

template <>

struct WrapBytes<LargeStringType> {

static inline PyObject* Wrap(const char* data, int64_t length) {

return PyUnicode_FromStringAndSize(data, length);

}

};

template <>

struct WrapBytes<StringViewType> {

static inline PyObject* Wrap(const char* data, int64_t length) {

return PyUnicode_FromStringAndSize(data, length);

}

};

template <>

struct WrapBytes<BinaryType> {

static inline PyObject* Wrap(const char* data, int64_t length) {

return PyBytes_FromStringAndSize(data, length);

}

};

template <>

struct WrapBytes<LargeBinaryType> {

static inline PyObject* Wrap(const char* data, int64_t length) {

return PyBytes_FromStringAndSize(data, length);

}

};

template <>

struct WrapBytes<BinaryViewType> {

static inline PyObject* Wrap(const char* data, int64_t length) {

return PyBytes_FromStringAndSize(data, length);

}

};

template <>

struct WrapBytes<FixedSizeBinaryType> {

static inline PyObject* Wrap(const char* data, int64_t length) {

return PyBytes_FromStringAndSize(data, length);

}

};

static inline bool ListTypeSupported(const DataType& type) {

switch (type.id()) {

case Type::BOOL:

case Type::UINT8:

case Type::INT8:

case Type::UINT16:

case Type::INT16:

case Type::UINT32:

case Type::INT32:

case Type::INT64:

case Type::UINT64:

case Type::HALF_FLOAT:

case Type::FLOAT:

case Type::DOUBLE:

case Type::DECIMAL128:

case Type::DECIMAL256:

case Type::BINARY:

case Type::LARGE_BINARY:

case Type::STRING:

case Type::LARGE_STRING:

case Type::DATE32:

case Type::DATE64:

case Type::STRUCT:

case Type::MAP:

case Type::TIME32:

case Type::TIME64:

case Type::TIMESTAMP:

case Type::DURATION:

case Type::DICTIONARY:

case Type::INTERVAL_MONTH_DAY_NANO:

case Type::NA: // empty list

// The above types are all supported.

return true;

case Type::FIXED_SIZE_LIST:

case Type::LIST:

case Type::LARGE_LIST:

case Type::LIST_VIEW:

case Type::LARGE_LIST_VIEW: {

const auto& list_type = checked_cast<const BaseListType&>(type);

return ListTypeSupported(*list_type.value_type());

}

case Type::EXTENSION: {

const auto& ext = checked_cast<const ExtensionType&>(*type.GetSharedPtr());

return ListTypeSupported(*(ext.storage_type()));

}

default:

break;

}

return false;

}

Status CapsulizeArray(const std::shared_ptr<Array>& arr, PyObject** out) {

auto capsule = new ArrayCapsule{{arr}};

*out = PyCapsule_New(reinterpret_cast<void*>(capsule), "arrow::Array",

&ArrayCapsule_Destructor);

if (*out == nullptr) {

delete capsule;

RETURN_IF_PYERROR();

}

return Status::OK();

}

Status CapsulizeBuffer(const std::shared_ptr<Buffer>& buffer, PyObject** out) {

auto capsule = new BufferCapsule{{buffer}};

*out = PyCapsule_New(reinterpret_cast<void*>(capsule), "arrow::Buffer",

&BufferCapsule_Destructor);

if (*out == nullptr) {

delete capsule;

RETURN_IF_PYERROR();

}

return Status::OK();

}

Status SetNdarrayBase(PyArrayObject* arr, PyObject* base) {

if (PyArray_SetBaseObject(arr, base) == -1) {

// Error occurred, trust that SetBaseObject sets the error state

Py_XDECREF(base);

RETURN_IF_PYERROR();

}

return Status::OK();

}

Status SetBufferBase(PyArrayObject* arr, const std::shared_ptr<Buffer>& buffer) {

PyObject* base;

RETURN_NOT_OK(CapsulizeBuffer(buffer, &base));

return SetNdarrayBase(arr, base);

}

inline void set_numpy_metadata(int type, const DataType* datatype, PyArray_Descr* out) {

auto metadata =

reinterpret_cast<PyArray_DatetimeDTypeMetaData*>(PyDataType_C_METADATA(out));

if (type == NPY_DATETIME) {

if (datatype->id() == Type::TIMESTAMP) {

const auto& timestamp_type = checked_cast<const TimestampType&>(*datatype);

metadata->meta.base = internal::NumPyFrequency(timestamp_type.unit());

} else {

ARROW_DCHECK(false)

<< "NPY_DATETIME views only supported for Arrow TIMESTAMP types";

}

} else if (type == NPY_TIMEDELTA) {

ARROW_DCHECK_EQ(datatype->id(), Type::DURATION);

const auto& duration_type = checked_cast<const DurationType&>(*datatype);

metadata->meta.base = internal::NumPyFrequency(duration_type.unit());

}

}

Status PyArray_NewFromPool(int nd, npy_intp* dims, PyArray_Descr* descr, MemoryPool* pool,

PyObject** out) {

// ARROW-6570: Allocate memory from MemoryPool for a couple reasons

//

// * Track allocations

// * Get better performance through custom allocators

int64_t total_size = PyDataType_ELSIZE(descr);

for (int i = 0; i < nd; ++i) {

total_size *= dims[i];

}

ARROW_ASSIGN_OR_RAISE(auto buffer, AllocateBuffer(total_size, pool));

*out = PyArray_NewFromDescr(&PyArray_Type, descr, nd, dims,

/*strides=*/nullptr,

/*data=*/buffer->mutable_data(),

/*flags=*/NPY_ARRAY_CARRAY | NPY_ARRAY_WRITEABLE,

/*obj=*/nullptr);

if (*out == nullptr) {

RETURN_IF_PYERROR();

// Trust that error set if NULL returned

}

return SetBufferBase(reinterpret_cast<PyArrayObject*>(*out), std::move(buffer));

}

template <typename T = void>

inline const T* GetPrimitiveValues(const Array& arr) {

if (arr.length() == 0) {

return nullptr;

}

const int elsize = arr.type()->byte_width();

const auto& prim_arr = checked_cast<const PrimitiveArray&>(arr);

return reinterpret_cast<const T*>(prim_arr.values()->data() + arr.offset() * elsize);

}

Status MakeNumPyView(std::shared_ptr<Array> arr, PyObject* py_ref, int npy_type, int ndim,

npy_intp* dims, PyObject** out) {

PyAcquireGIL lock;

PyArray_Descr* descr = internal::GetSafeNumPyDtype(npy_type);

set_numpy_metadata(npy_type, arr->type().get(), descr);

PyObject* result = PyArray_NewFromDescr(

&PyArray_Type, descr, ndim, dims, /*strides=*/nullptr,

const_cast<void*>(GetPrimitiveValues(*arr)), /*flags=*/0, nullptr);

PyArrayObject* np_arr = reinterpret_cast<PyArrayObject*>(result);

if (np_arr == nullptr) {

// Error occurred, trust that error set

return Status::OK();

}

PyObject* base;

if (py_ref == nullptr) {

// Capsule will be owned by the ndarray, no incref necessary. See

// ARROW-1973

RETURN_NOT_OK(CapsulizeArray(arr, &base));

} else {

Py_INCREF(py_ref);

base = py_ref;

}

RETURN_NOT_OK(SetNdarrayBase(np_arr, base));

// Do not allow Arrow data to be mutated

PyArray_CLEARFLAGS(np_arr, NPY_ARRAY_WRITEABLE);

*out = result;

return Status::OK();

}

class PandasWriter {

public:

enum type {

OBJECT,

UINT8,

INT8,

UINT16,

INT16,

UINT32,

INT32,

UINT64,

INT64,

HALF_FLOAT,

FLOAT,

DOUBLE,

BOOL,

DATETIME_DAY,

DATETIME_SECOND,

DATETIME_MILLI,

DATETIME_MICRO,

DATETIME_NANO,

DATETIME_SECOND_TZ,

DATETIME_MILLI_TZ,

DATETIME_MICRO_TZ,

DATETIME_NANO_TZ,

TIMEDELTA_SECOND,

TIMEDELTA_MILLI,

TIMEDELTA_MICRO,

TIMEDELTA_NANO,

CATEGORICAL,

EXTENSION

};

PandasWriter(const PandasOptions& options, int64_t num_rows, int num_columns)

: options_(options), num_rows_(num_rows), num_columns_(num_columns) {

PyAcquireGIL lock;

internal::InitPandasStaticData();

}

virtual ~PandasWriter() {}

void SetBlockData(PyObject* arr) {

block_arr_.reset(arr);

block_data_ =

reinterpret_cast<uint8_t*>(PyArray_DATA(reinterpret_cast<PyArrayObject*>(arr)));

}

/// \brief Either copy or wrap single array to create pandas-compatible array

/// for Series or DataFrame. num_columns_ can only be 1. Will try to zero

/// copy if possible (or error if not possible and zero_copy_only=True)

virtual Status TransferSingle(std::shared_ptr<ChunkedArray> data, PyObject* py_ref) = 0;

/// \brief Copy ChunkedArray into a multi-column block

virtual Status CopyInto(std::shared_ptr<ChunkedArray> data, int64_t rel_placement) = 0;

Status EnsurePlacementAllocated() {

std::lock_guard<std::mutex> guard(allocation_lock_);

if (placement_data_ != nullptr) {

return Status::OK();

}

PyAcquireGIL lock;

npy_intp placement_dims[1] = {num_columns_};

PyObject* placement_arr = PyArray_SimpleNew(1, placement_dims, NPY_INT64);

RETURN_IF_PYERROR();

placement_arr_.reset(placement_arr);

placement_data_ = reinterpret_cast<int64_t*>(

PyArray_DATA(reinterpret_cast<PyArrayObject*>(placement_arr)));

return Status::OK();

}

Status EnsureAllocated() {

std::lock_guard<std::mutex> guard(allocation_lock_);

if (block_data_ != nullptr) {

return Status::OK();

}

RETURN_NOT_OK(Allocate());

return Status::OK();

}

virtual bool CanZeroCopy(const ChunkedArray& data) const { return false; }

virtual Status Write(std::shared_ptr<ChunkedArray> data, int64_t abs_placement,

int64_t rel_placement) {

RETURN_NOT_OK(EnsurePlacementAllocated());

if (num_columns_ == 1 && options_.allow_zero_copy_blocks) {

RETURN_NOT_OK(TransferSingle(data, /*py_ref=*/nullptr));

} else {

RETURN_NOT_OK(

CheckNoZeroCopy("Cannot do zero copy conversion into "

"multi-column DataFrame block"));

RETURN_NOT_OK(EnsureAllocated());

RETURN_NOT_OK(CopyInto(data, rel_placement));

}

placement_data_[rel_placement] = abs_placement;

return Status::OK();

}

virtual Status GetDataFrameResult(PyObject** out) {

PyObject* result = PyDict_New();

RETURN_IF_PYERROR();

PyObject* block;

RETURN_NOT_OK(GetResultBlock(&block));

PyDict_SetItemString(result, "block", block);

PyDict_SetItemString(result, "placement", placement_arr_.obj());

RETURN_NOT_OK(AddResultMetadata(result));

*out = result;

return Status::OK();

}

// Caller steals the reference to this object

virtual Status GetSeriesResult(PyObject** out) {

RETURN_NOT_OK(MakeBlock1D());

// Caller owns the object now

*out = block_arr_.detach();

return Status::OK();

}

protected:

virtual Status AddResultMetadata(PyObject* result) { return Status::OK(); }

Status MakeBlock1D() {

// For Series or for certain DataFrame block types, we need to shape to a

// 1D array when there is only one column

PyAcquireGIL lock;

ARROW_DCHECK_EQ(1, num_columns_);

npy_intp new_dims[1] = {static_cast<npy_intp>(num_rows_)};

PyArray_Dims dims;

dims.ptr = new_dims;

dims.len = 1;

PyObject* reshaped = PyArray_Newshape(

reinterpret_cast<PyArrayObject*>(block_arr_.obj()), &dims, NPY_ANYORDER);

RETURN_IF_PYERROR();

// ARROW-8801: Here a PyArrayObject is created that is not being managed by

// any OwnedRef object. This object is then put in the resulting object

// with PyDict_SetItemString, which increments the reference count, so a

// memory leak ensues. There are several ways to fix the memory leak but a

// simple one is to put the reshaped 1D block array in this OwnedRefNoGIL

// so it will be correctly decref'd when this class is destructed.

block_arr_.reset(reshaped);

return Status::OK();

}

virtual Status GetResultBlock(PyObject** out) {

*out = block_arr_.obj();

return Status::OK();

}

Status CheckNoZeroCopy(const std::string& message) {

if (options_.zero_copy_only) {

return Status::Invalid(message);

}

return Status::OK();

}

Status CheckNotZeroCopyOnly(const ChunkedArray& data) {

if (options_.zero_copy_only) {

return Status::Invalid("Needed to copy ", data.num_chunks(), " chunks with ",

data.null_count(), " nulls, but zero_copy_only was True");

}

return Status::OK();

}

virtual Status Allocate() {

return Status::NotImplemented("Override Allocate in subclasses");

}

Status AllocateNDArray(int npy_type, int ndim = 2) {

PyAcquireGIL lock;

PyObject* block_arr = nullptr;

npy_intp block_dims[2] = {0, 0};

if (ndim == 2) {

block_dims[0] = num_columns_;

block_dims[1] = num_rows_;

} else {

block_dims[0] = num_rows_;

}

PyArray_Descr* descr = internal::GetSafeNumPyDtype(npy_type);

if (PyDataType_REFCHK(descr)) {

// ARROW-6876: if the array has refcounted items, let Numpy

// own the array memory so as to decref elements on array destruction

block_arr = PyArray_SimpleNewFromDescr(ndim, block_dims, descr);

RETURN_IF_PYERROR();

} else {

RETURN_NOT_OK(

PyArray_NewFromPool(ndim, block_dims, descr, options_.pool, &block_arr));

}

SetBlockData(block_arr);

return Status::OK();

}

void SetDatetimeUnit(NPY_DATETIMEUNIT unit) {

PyAcquireGIL lock;

auto date_dtype =

reinterpret_cast<PyArray_DatetimeDTypeMetaData*>(PyDataType_C_METADATA(

PyArray_DESCR(reinterpret_cast<PyArrayObject*>(block_arr_.obj()))));

date_dtype->meta.base = unit;

}

PandasOptions options_;

std::mutex allocation_lock_;

int64_t num_rows_;

int num_columns_;

OwnedRefNoGIL block_arr_;

uint8_t* block_data_ = nullptr;

// ndarray<int32>

OwnedRefNoGIL placement_arr_;

int64_t* placement_data_ = nullptr;

private:

ARROW_DISALLOW_COPY_AND_ASSIGN(PandasWriter);

};

template <typename InType, typename OutType>

inline void ConvertIntegerWithNulls(const PandasOptions& options,

const ChunkedArray& data, OutType* out_values) {

for (int c = 0; c < data.num_chunks(); c++) {

const auto& arr = *data.chunk(c);

const InType* in_values = GetPrimitiveValues<InType>(arr);

// Upcast to double, set NaN as appropriate

for (int i = 0; i < arr.length(); ++i) {

*out_values++ =

arr.IsNull(i) ? static_cast<OutType>(NAN) : static_cast<OutType>(in_values[i]);

}

}

}

template <typename T>

inline void ConvertIntegerNoNullsSameType(const PandasOptions& options,

const ChunkedArray& data, T* out_values) {

for (int c = 0; c < data.num_chunks(); c++) {

const auto& arr = *data.chunk(c);

if (arr.length() > 0) {

const T* in_values = GetPrimitiveValues<T>(arr);

memcpy(out_values, in_values, sizeof(T) * arr.length());

out_values += arr.length();

}

}

}

template <typename InType, typename OutType>

inline void ConvertIntegerNoNullsCast(const PandasOptions& options,

const ChunkedArray& data, OutType* out_values) {

for (int c = 0; c < data.num_chunks(); c++) {

const auto& arr = *data.chunk(c);

const InType* in_values = GetPrimitiveValues<InType>(arr);

for (int64_t i = 0; i < arr.length(); ++i) {

*out_values = in_values[i];

}

}

}

template <typename T, typename Enable = void>

struct MemoizationTraits {

using Scalar = typename T::c_type;

};

template <typename T>

struct MemoizationTraits<T, enable_if_has_string_view<T>> {

// For binary, we memoize string_view as a scalar value to avoid having to

// unnecessarily copy the memory into the memo table data structure

using Scalar = std::string_view;

};

// Generic Array -> PyObject** converter that handles object deduplication, if

// requested

template <typename Type, typename WrapFunction>

inline Status ConvertAsPyObjects(const PandasOptions& options, const ChunkedArray& data,

WrapFunction&& wrap_func, PyObject** out_values) {

using ArrayType = typename TypeTraits<Type>::ArrayType;

using Scalar = typename MemoizationTraits<Type>::Scalar;

auto convert_chunks = [&](auto&& wrap_func) -> Status {

for (int c = 0; c < data.num_chunks(); c++) {

const auto& arr = arrow::internal::checked_cast<const ArrayType&>(*data.chunk(c));

RETURN_NOT_OK(internal::WriteArrayObjects(arr, wrap_func, out_values));

out_values += arr.length();

}

return Status::OK();

};

if (options.deduplicate_objects) {

// GH-40316: only allocate a memo table if deduplication is enabled.

::arrow::internal::ScalarMemoTable<Scalar> memo_table(options.pool);

std::vector<PyObject*> unique_values;

int32_t memo_size = 0;

auto WrapMemoized = [&](const Scalar& value, PyObject** out_values) {

int32_t memo_index;

RETURN_NOT_OK(memo_table.GetOrInsert(value, &memo_index));

if (memo_index == memo_size) {

// New entry

RETURN_NOT_OK(wrap_func(value, out_values));

unique_values.push_back(*out_values);

++memo_size;

} else {

// Duplicate entry

Py_INCREF(unique_values[memo_index]);

*out_values = unique_values[memo_index];

}

return Status::OK();

};

return convert_chunks(std::move(WrapMemoized));

} else {

return convert_chunks(std::forward<WrapFunction>(wrap_func));

}

}

Status ConvertStruct(PandasOptions options, const ChunkedArray& data,

PyObject** out_values) {

if (data.num_chunks() == 0) {

return Status::OK();

}

// ChunkedArray has at least one chunk

auto arr = checked_cast<const StructArray*>(data.chunk(0).get());

// Use it to cache the struct type and number of fields for all chunks

int32_t num_fields = arr->num_fields();

auto array_type = arr->type();

std::vector<OwnedRef> fields_data(num_fields * data.num_chunks());

OwnedRef dict_item;

// See notes in MakeInnerOptions.

options = MakeInnerOptions(std::move(options));

// Don't blindly convert because timestamps in lists are handled differently.

options.timestamp_as_object = true;

for (int c = 0; c < data.num_chunks(); c++) {

auto fields_data_offset = c * num_fields;

auto arr = checked_cast<const StructArray*>(data.chunk(c).get());

// Convert the struct arrays first

for (int32_t i = 0; i < num_fields; i++) {

auto field = arr->field(static_cast<int>(i));

// In case the field is an extension array, use .storage() to convert to Pandas

if (field->type()->id() == Type::EXTENSION) {

const ExtensionArray& arr_ext = checked_cast<const ExtensionArray&>(*field);

field = arr_ext.storage();

}

RETURN_NOT_OK(ConvertArrayToPandas(options, field, nullptr,

fields_data[i + fields_data_offset].ref()));

ARROW_DCHECK(PyArray_Check(fields_data[i + fields_data_offset].obj()));

}

// Construct a dictionary for each row

const bool has_nulls = data.null_count() > 0;

for (int64_t i = 0; i < arr->length(); ++i) {

if (has_nulls && arr->IsNull(i)) {

Py_INCREF(Py_None);

*out_values = Py_None;

} else {

// Build the new dict object for the row

dict_item.reset(PyDict_New());

RETURN_IF_PYERROR();

for (int32_t field_idx = 0; field_idx < num_fields; ++field_idx) {

OwnedRef field_value;

auto name = array_type->field(static_cast<int>(field_idx))->name();

if (!arr->field(static_cast<int>(field_idx))->IsNull(i)) {

// Value exists in child array, obtain it

auto array = reinterpret_cast<PyArrayObject*>(

fields_data[field_idx + fields_data_offset].obj());

auto ptr = reinterpret_cast<const char*>(PyArray_GETPTR1(array, i));

field_value.reset(PyArray_GETITEM(array, ptr));

RETURN_IF_PYERROR();

} else {

// Translate the Null to a None

Py_INCREF(Py_None);

field_value.reset(Py_None);

}

// PyDict_SetItemString increments reference count

auto setitem_result =

PyDict_SetItemString(dict_item.obj(), name.c_str(), field_value.obj());

RETURN_IF_PYERROR();

ARROW_DCHECK_EQ(setitem_result, 0);

}

*out_values = dict_item.obj();

// Grant ownership to the resulting array

Py_INCREF(*out_values);

}

++out_values;

}

}

return Status::OK();

}

Status DecodeDictionaries(MemoryPool* pool, const std::shared_ptr<DataType>& dense_type,

ArrayVector* arrays) {

compute::ExecContext ctx(pool);

compute::CastOptions options;

for (size_t i = 0; i < arrays->size(); ++i) {

ARROW_ASSIGN_OR_RAISE((*arrays)[i],

compute::Cast(*(*arrays)[i], dense_type, options, &ctx));

}

return Status::OK();

}

Status DecodeDictionaries(MemoryPool* pool, const std::shared_ptr<DataType>& dense_type,

std::shared_ptr<ChunkedArray>* array) {

auto chunks = (*array)->chunks();

RETURN_NOT_OK(DecodeDictionaries(pool, dense_type, &chunks));

*array = std::make_shared<ChunkedArray>(std::move(chunks), dense_type);

return Status::OK();

}

template <typename T>

enable_if_list_like<T, Status> ConvertListsLike(PandasOptions options,

const ChunkedArray& data,

PyObject** out_values) {

using ListArrayT = typename TypeTraits<T>::ArrayType;

// Get column of underlying value arrays

ArrayVector value_arrays;

for (int c = 0; c < data.num_chunks(); c++) {

const auto& arr = checked_cast<const ListArrayT&>(*data.chunk(c));

// values() does not account for offsets, so we need to slice into it.

// We can't use Flatten(), because it removes the values behind a null list

// value, and that makes the offsets into original list values and our

// flattened_values array different.

std::shared_ptr<Array> flattened_values = arr.values()->Slice(

arr.value_offset(0), arr.value_offset(arr.length()) - arr.value_offset(0));

if (arr.value_type()->id() == Type::EXTENSION) {

const auto& arr_ext = checked_cast<const ExtensionArray&>(*flattened_values);

value_arrays.emplace_back(arr_ext.storage());

} else {

value_arrays.emplace_back(flattened_values);

}

}

using ListArrayType = typename ListArrayT::TypeClass;

const auto& list_type = checked_cast<const ListArrayType&>(*data.type());

auto value_type = list_type.value_type();

if (value_type->id() == Type::EXTENSION) {

value_type = checked_cast<const ExtensionType&>(*value_type).storage_type();

}

auto flat_column = std::make_shared<ChunkedArray>(value_arrays, value_type);

options = MakeInnerOptions(std::move(options));

OwnedRefNoGIL owned_numpy_array;

RETURN_NOT_OK(ConvertChunkedArrayToPandas(options, flat_column, nullptr,

owned_numpy_array.ref()));

PyObject* numpy_array = owned_numpy_array.obj();

ARROW_DCHECK(PyArray_Check(numpy_array));

int64_t chunk_offset = 0;

for (int c = 0; c < data.num_chunks(); c++) {

const auto& arr = checked_cast<const ListArrayT&>(*data.chunk(c));

const bool has_nulls = data.null_count() > 0;

for (int64_t i = 0; i < arr.length(); ++i) {

if (has_nulls && arr.IsNull(i)) {

Py_INCREF(Py_None);

*out_values = Py_None;

} else {

// Need to subtract value_offset(0) since the original chunk might be a slice

// into another array.

OwnedRef start(PyLong_FromLongLong(arr.value_offset(i) + chunk_offset -

arr.value_offset(0)));

OwnedRef end(PyLong_FromLongLong(arr.value_offset(i + 1) + chunk_offset -

arr.value_offset(0)));

OwnedRef slice(PySlice_New(start.obj(), end.obj(), nullptr));

if (ARROW_PREDICT_FALSE(slice.obj() == nullptr)) {

// Fall out of loop, will return from RETURN_IF_PYERROR

break;

}

*out_values = PyObject_GetItem(numpy_array, slice.obj());

if (*out_values == nullptr) {

// Fall out of loop, will return from RETURN_IF_PYERROR

break;

}

}

++out_values;

}

RETURN_IF_PYERROR();

chunk_offset += arr.value_offset(arr.length()) - arr.value_offset(0);

}

return Status::OK();

}

// TODO GH-40579: optimize ListView conversion to avoid unnecessary copies

template <typename T>

enable_if_list_view<T, Status> ConvertListsLike(PandasOptions options,

const ChunkedArray& data,

PyObject** out_values) {

using ListViewArrayType = typename TypeTraits<T>::ArrayType;

using NonViewType =

std::conditional_t<T::type_id == Type::LIST_VIEW, ListType, LargeListType>;

using NonViewClass = typename TypeTraits<NonViewType>::ArrayType;

ArrayVector list_arrays;

for (int c = 0; c < data.num_chunks(); c++) {

const auto& arr = checked_cast<const ListViewArrayType&>(*data.chunk(c));

ARROW_ASSIGN_OR_RAISE(auto non_view_array,

NonViewClass::FromListView(arr, options.pool));

list_arrays.emplace_back(non_view_array);

}

auto chunked_array = std::make_shared<ChunkedArray>(list_arrays);

return ConvertListsLike<NonViewType>(options, *chunked_array, out_values);

}

template <typename F1, typename F2, typename F3>

Status ConvertMapHelper(F1 resetRow, F2 addPairToRow, F3 stealRow,

const ChunkedArray& data, PyArrayObject* py_keys,

PyArrayObject* py_items,

// needed for null checks in items

const std::vector<std::shared_ptr<Array>> item_arrays,

PyObject** out_values) {

OwnedRef key_value;

OwnedRef item_value;

int64_t chunk_offset = 0;

for (int c = 0; c < data.num_chunks(); ++c) {

const auto& arr = checked_cast<const MapArray&>(*data.chunk(c));

const bool has_nulls = data.null_count() > 0;

// Make a list of key/item pairs for each row in array

for (int64_t i = 0; i < arr.length(); ++i) {

if (has_nulls && arr.IsNull(i)) {

Py_INCREF(Py_None);

*out_values = Py_None;

} else {

int64_t entry_offset = arr.value_offset(i);

int64_t num_pairs = arr.value_offset(i + 1) - entry_offset;

// Build the new list object for the row of Python pairs

RETURN_NOT_OK(resetRow(num_pairs));

// Add each key/item pair in the row

for (int64_t j = 0; j < num_pairs; ++j) {

// Get key value, key is non-nullable for a valid row

auto ptr_key = reinterpret_cast<const char*>(

PyArray_GETPTR1(py_keys, chunk_offset + entry_offset + j));

key_value.reset(PyArray_GETITEM(py_keys, ptr_key));

RETURN_IF_PYERROR();

if (item_arrays[c]->IsNull(entry_offset + j)) {

// Translate the Null to a None

Py_INCREF(Py_None);

item_value.reset(Py_None);

} else {

// Get valid value from item array

auto ptr_item = reinterpret_cast<const char*>(

PyArray_GETPTR1(py_items, chunk_offset + entry_offset + j));

item_value.reset(PyArray_GETITEM(py_items, ptr_item));

RETURN_IF_PYERROR();

}

// Add the key/item pair to the row

RETURN_NOT_OK(addPairToRow(j, key_value, item_value));

}

// Pass ownership to the resulting array

*out_values = stealRow();

}

++out_values;

}

RETURN_IF_PYERROR();

chunk_offset += arr.values()->length();

}

return Status::OK();

}

// A more helpful error message around TypeErrors that may stem from unhashable keys

Status CheckMapAsPydictsTypeError() {

if (ARROW_PREDICT_TRUE(!PyErr_Occurred())) {

return Status::OK();

}

if (PyErr_ExceptionMatches(PyExc_TypeError)) {

// Modify the error string directly, so it is re-raised

// with our additional info.

//

// There are not many interesting things happening when this

// is hit. This is intended to only be called directly after

// PyDict_SetItem, where a finite set of errors could occur.

PyObject *type, *value, *traceback;

PyErr_Fetch(&type, &value, &traceback);

std::string message;

RETURN_NOT_OK(internal::PyObject_StdStringStr(value, &message));

message +=

". If keys are not hashable, then you must use the option "

"[maps_as_pydicts=None (default)]";

// resets the error

PyErr_SetString(PyExc_TypeError, message.c_str());

}

return ConvertPyError();

}

Status CheckForDuplicateKeys(bool error_on_duplicate_keys, Py_ssize_t total_dict_len,

Py_ssize_t total_raw_len) {

if (total_dict_len < total_raw_len) {

const char* message =

"[maps_as_pydicts] "

"After conversion of Arrow maps to pydicts, "

"detected data loss due to duplicate keys. "

"Original input length is [%lld], total converted pydict length is [%lld].";

std::array<char, 256> buf;

std::snprintf(buf.data(), buf.size(), message, total_raw_len, total_dict_len);

if (error_on_duplicate_keys) {

return Status::UnknownError(buf.data());

} else {

ARROW_LOG(WARNING) << buf.data();

}

}

return Status::OK();

}

Status ConvertMap(PandasOptions options, const ChunkedArray& data,

PyObject** out_values) {

// Get columns of underlying key/item arrays

std::vector<std::shared_ptr<Array>> key_arrays;

std::vector<std::shared_ptr<Array>> item_arrays;

for (int c = 0; c < data.num_chunks(); ++c) {

const auto& map_arr = checked_cast<const MapArray&>(*data.chunk(c));

key_arrays.emplace_back(map_arr.keys());

item_arrays.emplace_back(map_arr.items());

}

const auto& map_type = checked_cast<const MapType&>(*data.type());

auto key_type = map_type.key_type();

auto item_type = map_type.item_type();

// ARROW-6899: Convert dictionary-encoded children to dense instead of

// failing below. A more efficient conversion than this could be done later

if (key_type->id() == Type::DICTIONARY) {

auto dense_type = checked_cast<const DictionaryType&>(*key_type).value_type();

RETURN_NOT_OK(DecodeDictionaries(options.pool, dense_type, &key_arrays));

key_type = dense_type;

}

if (item_type->id() == Type::DICTIONARY) {

auto dense_type = checked_cast<const DictionaryType&>(*item_type).value_type();