Skip to content

Commit

Permalink
Initial commit
Browse files Browse the repository at this point in the history
  • Loading branch information
@AlenkaF
AlenkaF committed May 21, 2024
1 parent e3cd0ae commit 27f3ae2
Show file tree
Hide file tree
Showing 4 changed files with 33 additions and 22 deletions.
27 changes: 13 additions & 14 deletions cpp/src/arrow/array/array_nested.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -807,7 +807,7 @@ MapArray::MapArray(const std::shared_ptr<DataType>& type, int64_t length,
Result<std::shared_ptr<Array>> MapArray::FromArraysInternal(
std::shared_ptr<DataType> type, const std::shared_ptr<Array>& offsets,
const std::shared_ptr<Array>& keys, const std::shared_ptr<Array>& items,
MemoryPool* pool) {
MemoryPool* pool, const std::shared_ptr<Buffer>& null_bitmap) {
using offset_type = typename MapType::offset_type;
using OffsetArrowType = typename CTypeTraits<offset_type>::ArrowType;

Expand All @@ -829,31 +829,30 @@ Result<std::shared_ptr<Array>> MapArray::FromArraysInternal(

if (offsets->null_count() > 0) {
ARROW_ASSIGN_OR_RAISE(auto buffers,
CleanListOffsets<MapType>(NULLPTR, *offsets, pool));
CleanListOffsets<MapType>(null_bitmap, *offsets, pool));
return std::make_shared<MapArray>(type, offsets->length() - 1, std::move(buffers),
keys, items, offsets->null_count(), 0);
}

using OffsetArrayType = typename TypeTraits<OffsetArrowType>::ArrayType;
const auto& typed_offsets = checked_cast<const OffsetArrayType&>(*offsets);
auto buffers = BufferVector({nullptr, typed_offsets.values()});
auto buffers = BufferVector({null_bitmap, typed_offsets.values()});
return std::make_shared<MapArray>(type, offsets->length() - 1, std::move(buffers), keys,
items, /*null_count=*/0, offsets->offset());
}

Result<std::shared_ptr<Array>> MapArray::FromArrays(const std::shared_ptr<Array>& offsets,
const std::shared_ptr<Array>& keys,
const std::shared_ptr<Array>& items,
MemoryPool* pool) {
Result<std::shared_ptr<Array>> MapArray::FromArrays(
const std::shared_ptr<Array>& offsets, const std::shared_ptr<Array>& keys,
const std::shared_ptr<Array>& items, MemoryPool* pool,
const std::shared_ptr<Buffer>& null_bitmap) {
return FromArraysInternal(std::make_shared<MapType>(keys->type(), items->type()),
offsets, keys, items, pool);
offsets, keys, items, pool, null_bitmap);
}

Result<std::shared_ptr<Array>> MapArray::FromArrays(std::shared_ptr<DataType> type,
const std::shared_ptr<Array>& offsets,
const std::shared_ptr<Array>& keys,
const std::shared_ptr<Array>& items,
MemoryPool* pool) {
Result<std::shared_ptr<Array>> MapArray::FromArrays(
std::shared_ptr<DataType> type, const std::shared_ptr<Array>& offsets,
const std::shared_ptr<Array>& keys, const std::shared_ptr<Array>& items,
MemoryPool* pool, const std::shared_ptr<Buffer>& null_bitmap) {
if (type->id() != Type::MAP) {
return Status::TypeError("Expected map type, got ", type->ToString());
}
Expand All @@ -864,7 +863,7 @@ Result<std::shared_ptr<Array>> MapArray::FromArrays(std::shared_ptr<DataType> ty
if (!map_type.item_type()->Equals(items->type())) {
return Status::TypeError("Mismatching map items type");
}
return FromArraysInternal(std::move(type), offsets, keys, items, pool);
return FromArraysInternal(std::move(type), offsets, keys, items, pool, null_bitmap);
}

Status MapArray::ValidateChildData(
Expand Down
9 changes: 6 additions & 3 deletions cpp/src/arrow/array/array_nested.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -532,15 +532,18 @@ class ARROW_EXPORT MapArray : public ListArray {
/// \param[in] keys Array containing key values
/// \param[in] items Array containing item values
/// \param[in] pool MemoryPool in case new offsets array needs to be
/// \param[in] null_bitmap Optional validity bitmap
/// allocated because of null values
static Result<std::shared_ptr<Array>> FromArrays(
const std::shared_ptr<Array>& offsets, const std::shared_ptr<Array>& keys,
const std::shared_ptr<Array>& items, MemoryPool* pool = default_memory_pool());
const std::shared_ptr<Array>& items, MemoryPool* pool = default_memory_pool(),
const std::shared_ptr<Buffer>& null_bitmap = NULLPTR);

static Result<std::shared_ptr<Array>> FromArrays(
std::shared_ptr<DataType> type, const std::shared_ptr<Array>& offsets,
const std::shared_ptr<Array>& keys, const std::shared_ptr<Array>& items,
MemoryPool* pool = default_memory_pool());
MemoryPool* pool = default_memory_pool(),
const std::shared_ptr<Buffer>& null_bitmap = NULLPTR);

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

Expand All @@ -560,7 +563,7 @@ class ARROW_EXPORT MapArray : public ListArray {
static Result<std::shared_ptr<Array>> FromArraysInternal(
std::shared_ptr<DataType> type, const std::shared_ptr<Array>& offsets,
const std::shared_ptr<Array>& keys, const std::shared_ptr<Array>& items,
MemoryPool* pool);
MemoryPool* pool, const std::shared_ptr<Buffer>& null_bitmap = NULLPTR);

private:
const MapType* map_type_;
Expand Down
11 changes: 8 additions & 3 deletions python/pyarrow/array.pxi
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3060,7 +3060,7 @@ cdef class MapArray(ListArray):
"""

@staticmethod
def from_arrays(offsets, keys, items, DataType type=None, MemoryPool pool=None):
def from_arrays(offsets, keys, items, DataType type=None, MemoryPool pool=None, mask=None):
"""
Construct MapArray from arrays of int32 offsets and key, item arrays.
Expand All @@ -3072,6 +3072,8 @@ cdef class MapArray(ListArray):
type : DataType, optional
If not specified, a default MapArray with the keys' and items' type is used.
pool : MemoryPool
mask : Array (boolean type), optional
Indicate which values are null (True) or not null (False).
Returns
-------
Expand Down Expand Up @@ -3153,24 +3155,27 @@ cdef class MapArray(ListArray):
cdef:
Array _offsets, _keys, _items
shared_ptr[CArray] out
shared_ptr[CBuffer] c_mask
cdef CMemoryPool* cpool = maybe_unbox_memory_pool(pool)

_offsets = asarray(offsets, type='int32')
_keys = asarray(keys)
_items = asarray(items)

c_mask = c_mask_inverted_from_obj(mask, pool)

if type is not None:
with nogil:
out = GetResultValue(
CMapArray.FromArraysAndType(
type.sp_type, _offsets.sp_array,
_keys.sp_array, _items.sp_array, cpool))
_keys.sp_array, _items.sp_array, cpool, c_mask))
else:
with nogil:
out = GetResultValue(
CMapArray.FromArrays(_offsets.sp_array,
_keys.sp_array,
_items.sp_array, cpool))
_items.sp_array, cpool, c_mask))
cdef Array result = pyarrow_wrap_array(out)
result.validate()
return result
Expand Down
8 changes: 6 additions & 2 deletions python/pyarrow/includes/libarrow.pxd
View file
Original file line number Diff line number Diff line change
Expand Up @@ -788,15 +788,19 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
const shared_ptr[CArray]& offsets,
const shared_ptr[CArray]& keys,
const shared_ptr[CArray]& items,
CMemoryPool* pool)
CMemoryPool* pool,
const shared_ptr[CBuffer] null_bitmap,
)

@staticmethod
CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"(
shared_ptr[CDataType],
const shared_ptr[CArray]& offsets,
const shared_ptr[CArray]& keys,
const shared_ptr[CArray]& items,
CMemoryPool* pool)
CMemoryPool* pool,
const shared_ptr[CBuffer] null_bitmap,
)

shared_ptr[CArray] keys()
shared_ptr[CArray] items()
Expand Down

0 comments on commit 27f3ae2

Please sign in to comment.