Skip to content

Commit 91aeae8

Browse files
milesgrangerpitrou
milesgranger
and
pitrou
authored
ARROW-15839: [C++][Python] Accept validity bitmap in ListArray.from_arrays (apache#13894)
Will close [ARROW-15839](https://issues.apache.org/jira/browse/ARROW-15839) This will allow proper nulls in place of what is presently empty lists (on top level) when no mask is passed. ```python import pyarrow as pa arr = pa.array([None, [0]]) reconstructed_arr = pa.ListArray.from_arrays(arr.offsets, arr.values, mask=arr.is_null()) print(reconstructed_arr.to_pylist()) # [None, [0]] # Currently it doesn't accept `mask` and results in `[[], [0]]` unless explicit offsets set ``` There was also discussion with @jorisvandenbossche about renaming the added `null_bitmap` to `validity_bitmap`; but this will deviate from [existing `null_bitmap` use in the same area of code](https://github.com/apache/arrow/blob/d880d7517a33f2ac8ff259cad711bc210fd570c5/cpp/src/arrow/array/array_nested.h#L113). Should I change those names now or in a later PR/issue to standardize the `null_bitmap` / `validity_buf` naming? Lead-authored-by: Miles Granger <miles59923@gmail.com> Co-authored-by: Antoine Pitrou <pitrou@free.fr> Signed-off-by: Antoine Pitrou <antoine@python.org>
1 parent d3d6371 commit 91aeae8

6 files changed

Lines changed: 209 additions & 53 deletions

File tree

‎cpp/src/arrow/array/array_list_test.cc‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,42 @@ class TestListArray : public ::testing::Test {
197197
EXPECT_FALSE(left->Slice(offset)->Equals(right->Slice(offset)));
198198
}
199199

200+
void TestFromArraysWithNullBitMap() {
201+
std::shared_ptr<Array> offsets_w_nulls, offsets_wo_nulls, values;
202+
203+
std::vector<offset_type> offsets = {0, 1, 1, 3, 4};
204+
std::vector<bool> offsets_w_nulls_is_valid = {true, false, true, true, true};
205+
206+
ArrayFromVector<OffsetType, offset_type>(offsets_w_nulls_is_valid, offsets,
207+
&offsets_w_nulls);
208+
ArrayFromVector<OffsetType, offset_type>(offsets, &offsets_wo_nulls);
209+
210+
auto type = std::make_shared<T>(int32());
211+
auto expected = std::dynamic_pointer_cast<ArrayType>(
212+
ArrayFromJSON(type, "[[0], null, [0, null], [0]]"));
213+
values = expected->values();
214+
215+
// Offsets with nulls will match.
216+
ASSERT_OK_AND_ASSIGN(auto result,
217+
ArrayType::FromArrays(*offsets_w_nulls, *values, pool_));
218+
AssertArraysEqual(*result, *expected);
219+
220+
// Offets without nulls, will replace null with empty list
221+
ASSERT_OK_AND_ASSIGN(result,
222+
ArrayType::FromArrays(*offsets_wo_nulls, *values, pool_));
223+
AssertArraysEqual(*result, *std::dynamic_pointer_cast<ArrayType>(
224+
ArrayFromJSON(type, "[[0], [], [0, null], [0]]")));
225+
226+
// Specify non-null offsets with null_bitmap
227+
ASSERT_OK_AND_ASSIGN(result, ArrayType::FromArrays(*offsets_wo_nulls, *values, pool_,
228+
expected->null_bitmap()));
229+
AssertArraysEqual(*result, *expected);
230+
231+
// Cannot specify both null offsets with null_bitmap
232+
ASSERT_RAISES(Invalid, ArrayType::FromArrays(*offsets_w_nulls, *values, pool_,
233+
expected->null_bitmap()));
234+
}
235+
200236
void TestFromArrays() {
201237
std::shared_ptr<Array> offsets1, offsets2, offsets3, offsets4, offsets5, values;
202238

@@ -539,6 +575,10 @@ TYPED_TEST(TestListArray, ValuesEquality) { this->TestValuesEquality(); }
539575

540576
TYPED_TEST(TestListArray, FromArrays) { this->TestFromArrays(); }
541577

578+
TYPED_TEST(TestListArray, FromArraysWithNullBitMap) {
579+
this->TestFromArraysWithNullBitMap();
580+
}
581+
542582
TYPED_TEST(TestListArray, AppendNull) { this->TestAppendNull(); }
543583

544584
TYPED_TEST(TestListArray, AppendNulls) { this->TestAppendNulls(); }

‎cpp/src/arrow/array/array_nested.cc‎

Lines changed: 33 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -103,7 +103,8 @@ Status CleanListOffsets(const Array& offsets, MemoryPool* pool,
103103
template <typename TYPE>
104104
Result<std::shared_ptr<typename TypeTraits<TYPE>::ArrayType>> ListArrayFromArrays(
105105
std::shared_ptr<DataType> type, const Array& offsets, const Array& values,
106-
MemoryPool* pool) {
106+
MemoryPool* pool, std::shared_ptr<Buffer> null_bitmap = NULLPTR,
107+
int64_t null_count = kUnknownNullCount) {
107108
using offset_type = typename TYPE::offset_type;
108109
using ArrayType = typename TypeTraits<TYPE>::ArrayType;
109110
using OffsetArrowType = typename CTypeTraits<offset_type>::ArrowType;
@@ -116,14 +117,24 @@ Result<std::shared_ptr<typename TypeTraits<TYPE>::ArrayType>> ListArrayFromArray
116117
return Status::TypeError("List offsets must be ", OffsetArrowType::type_name());
117118
}
118119

120+
if (null_bitmap != nullptr && offsets.null_count() > 0) {
121+
return Status::Invalid(
122+
"Ambiguous to specify both validity map and offsets with nulls");
123+
}
124+
125+
if (null_bitmap != nullptr && offsets.offset() != 0) {
126+
return Status::NotImplemented("Null bitmap with offsets slice not supported.");
127+
}
128+
119129
std::shared_ptr<Buffer> offset_buf, validity_buf;
120130
RETURN_NOT_OK(CleanListOffsets<TYPE>(offsets, pool, &offset_buf, &validity_buf));
121-
BufferVector buffers = {validity_buf, offset_buf};
131+
int64_t null_count_ = null_bitmap ? null_count : offsets.null_count();
132+
BufferVector buffers = {null_bitmap ? std::move(null_bitmap) : validity_buf,
133+
offset_buf};
122134

123-
auto internal_data = ArrayData::Make(type, offsets.length() - 1, std::move(buffers),
124-
offsets.null_count(), offsets.offset());
135+
std::shared_ptr<arrow::ArrayData> internal_data = ArrayData::Make(
136+
type, offsets.length() - 1, std::move(buffers), null_count_, offsets.offset());
125137
internal_data->child_data.push_back(values.data());
126-
127138
return std::make_shared<ArrayType>(internal_data);
128139
}
129140

@@ -231,45 +242,47 @@ void LargeListArray::SetData(const std::shared_ptr<ArrayData>& data) {
231242
internal::SetListData(this, data);
232243
}
233244

234-
Result<std::shared_ptr<ListArray>> ListArray::FromArrays(const Array& offsets,
235-
const Array& values,
236-
MemoryPool* pool) {
245+
Result<std::shared_ptr<ListArray>> ListArray::FromArrays(
246+
const Array& offsets, const Array& values, MemoryPool* pool,
247+
std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
237248
return ListArrayFromArrays<ListType>(std::make_shared<ListType>(values.type()), offsets,
238-
values, pool);
249+
values, pool, null_bitmap, null_count);
239250
}
240251

241-
Result<std::shared_ptr<ListArray>> ListArray::FromArrays(std::shared_ptr<DataType> type,
242-
const Array& offsets,
243-
const Array& values,
244-
MemoryPool* pool) {
252+
Result<std::shared_ptr<ListArray>> ListArray::FromArrays(
253+
std::shared_ptr<DataType> type, const Array& offsets, const Array& values,
254+
MemoryPool* pool, std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
245255
if (type->id() != Type::LIST) {
246256
return Status::TypeError("Expected list type, got ", type->ToString());
247257
}
248258
const auto& list_type = checked_cast<const ListType&>(*type);
249259
if (!list_type.value_type()->Equals(values.type())) {
250260
return Status::TypeError("Mismatching list value type");
251261
}
252-
return ListArrayFromArrays<ListType>(std::move(type), offsets, values, pool);
262+
return ListArrayFromArrays<ListType>(std::move(type), offsets, values, pool,
263+
null_bitmap, null_count);
253264
}
254265

255-
Result<std::shared_ptr<LargeListArray>> LargeListArray::FromArrays(const Array& offsets,
256-
const Array& values,
257-
MemoryPool* pool) {
266+
Result<std::shared_ptr<LargeListArray>> LargeListArray::FromArrays(
267+
const Array& offsets, const Array& values, MemoryPool* pool,
268+
std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
258269
return ListArrayFromArrays<LargeListType>(
259-
std::make_shared<LargeListType>(values.type()), offsets, values, pool);
270+
std::make_shared<LargeListType>(values.type()), offsets, values, pool, null_bitmap,
271+
null_count);
260272
}
261273

262274
Result<std::shared_ptr<LargeListArray>> LargeListArray::FromArrays(
263275
std::shared_ptr<DataType> type, const Array& offsets, const Array& values,
264-
MemoryPool* pool) {
276+
MemoryPool* pool, std::shared_ptr<Buffer> null_bitmap, int64_t null_count) {
265277
if (type->id() != Type::LARGE_LIST) {
266278
return Status::TypeError("Expected large list type, got ", type->ToString());
267279
}
268280
const auto& list_type = checked_cast<const LargeListType&>(*type);
269281
if (!list_type.value_type()->Equals(values.type())) {
270282
return Status::TypeError("Mismatching list value type");
271283
}
272-
return ListArrayFromArrays<LargeListType>(std::move(type), offsets, values, pool);
284+
return ListArrayFromArrays<LargeListType>(std::move(type), offsets, values, pool,
285+
null_bitmap, null_count);
273286
}
274287

275288
Result<std::shared_ptr<Array>> ListArray::Flatten(MemoryPool* memory_pool) const {

‎cpp/src/arrow/array/array_nested.h‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -120,18 +120,26 @@ class ARROW_EXPORT ListArray : public BaseListArray<ListType> {
120120
/// the offsets contain any nulls). If the offsets do not have nulls, they
121121
/// are assumed to be well-formed
122122
///
123+
/// Offsets of an Array's null bitmap can be present or an explicit
124+
/// null_bitmap, but not both.
125+
///
123126
/// \param[in] offsets Array containing n + 1 offsets encoding length and
124127
/// size. Must be of int32 type
125128
/// \param[in] values Array containing list values
126129
/// \param[in] pool MemoryPool in case new offsets array needs to be
127130
/// allocated because of null values
131+
/// \param[in] null_bitmap Optional validity bitmap
132+
/// \param[in] null_count Optional null count in null_bitmap
128133
static Result<std::shared_ptr<ListArray>> FromArrays(
129-
const Array& offsets, const Array& values,
130-
MemoryPool* pool = default_memory_pool());
134+
const Array& offsets, const Array& values, MemoryPool* pool = default_memory_pool(),
135+
std::shared_ptr<Buffer> null_bitmap = NULLPTR,
136+
int64_t null_count = kUnknownNullCount);
131137

132138
static Result<std::shared_ptr<ListArray>> FromArrays(
133139
std::shared_ptr<DataType> type, const Array& offsets, const Array& values,
134-
MemoryPool* pool = default_memory_pool());
140+
MemoryPool* pool = default_memory_pool(),
141+
std::shared_ptr<Buffer> null_bitmap = NULLPTR,
142+
int64_t null_count = kUnknownNullCount);
135143

136144
/// \brief Return an Array that is a concatenation of the lists in this array.
137145
///
@@ -178,13 +186,18 @@ class ARROW_EXPORT LargeListArray : public BaseListArray<LargeListType> {
178186
/// \param[in] values Array containing list values
179187
/// \param[in] pool MemoryPool in case new offsets array needs to be
180188
/// allocated because of null values
189+
/// \param[in] null_bitmap Optional validity bitmap
190+
/// \param[in] null_count Optional null count in null_bitmap
181191
static Result<std::shared_ptr<LargeListArray>> FromArrays(
182-
const Array& offsets, const Array& values,
183-
MemoryPool* pool = default_memory_pool());
192+
const Array& offsets, const Array& values, MemoryPool* pool = default_memory_pool(),
193+
std::shared_ptr<Buffer> null_bitmap = NULLPTR,
194+
int64_t null_count = kUnknownNullCount);
184195

185196
static Result<std::shared_ptr<LargeListArray>> FromArrays(
186197
std::shared_ptr<DataType> type, const Array& offsets, const Array& values,
187-
MemoryPool* pool = default_memory_pool());
198+
MemoryPool* pool = default_memory_pool(),
199+
std::shared_ptr<Buffer> null_bitmap = NULLPTR,
200+
int64_t null_count = kUnknownNullCount);
188201

189202
/// \brief Return an Array that is a concatenation of the lists in this array.
190203
///

‎python/pyarrow/array.pxi‎

Lines changed: 39 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1890,7 +1890,7 @@ cdef class ListArray(BaseListArray):
18901890
"""
18911891

18921892
@staticmethod
1893-
def from_arrays(offsets, values, DataType type=None, MemoryPool pool=None):
1893+
def from_arrays(offsets, values, DataType type=None, MemoryPool pool=None, mask=None):
18941894
"""
18951895
Construct ListArray from arrays of int32 offsets and values.
18961896
@@ -1901,7 +1901,9 @@ cdef class ListArray(BaseListArray):
19011901
type : DataType, optional
19021902
If not specified, a default ListType with the values' type is
19031903
used.
1904-
pool : MemoryPool
1904+
pool : MemoryPool, optional
1905+
mask : Array (boolean type), optional
1906+
Indicate which values are null (True) or not null (False).
19051907
19061908
Returns
19071909
-------
@@ -1943,21 +1945,24 @@ cdef class ListArray(BaseListArray):
19431945
cdef:
19441946
Array _offsets, _values
19451947
shared_ptr[CArray] out
1948+
shared_ptr[CBuffer] c_mask
19461949
cdef CMemoryPool* cpool = maybe_unbox_memory_pool(pool)
19471950

19481951
_offsets = asarray(offsets, type='int32')
19491952
_values = asarray(values)
19501953

1954+
c_mask = c_mask_inverted_from_obj(mask, pool)
1955+
19511956
if type is not None:
19521957
with nogil:
19531958
out = GetResultValue(
19541959
CListArray.FromArraysAndType(
1955-
type.sp_type, _offsets.ap[0], _values.ap[0], cpool))
1960+
type.sp_type, _offsets.ap[0], _values.ap[0], cpool, c_mask))
19561961
else:
19571962
with nogil:
19581963
out = GetResultValue(
19591964
CListArray.FromArrays(
1960-
_offsets.ap[0], _values.ap[0], cpool))
1965+
_offsets.ap[0], _values.ap[0], cpool, c_mask))
19611966
cdef Array result = pyarrow_wrap_array(out)
19621967
result.validate()
19631968
return result
@@ -2004,7 +2009,7 @@ cdef class LargeListArray(BaseListArray):
20042009
"""
20052010

20062011
@staticmethod
2007-
def from_arrays(offsets, values, DataType type=None, MemoryPool pool=None):
2012+
def from_arrays(offsets, values, DataType type=None, MemoryPool pool=None, mask=None):
20082013
"""
20092014
Construct LargeListArray from arrays of int64 offsets and values.
20102015
@@ -2015,7 +2020,9 @@ cdef class LargeListArray(BaseListArray):
20152020
type : DataType, optional
20162021
If not specified, a default ListType with the values' type is
20172022
used.
2018-
pool : MemoryPool
2023+
pool : MemoryPool, optional
2024+
mask : Array (boolean type), optional
2025+
Indicate which values are null (True) or not null (False).
20192026
20202027
Returns
20212028
-------
@@ -2024,21 +2031,25 @@ cdef class LargeListArray(BaseListArray):
20242031
cdef:
20252032
Array _offsets, _values
20262033
shared_ptr[CArray] out
2034+
shared_ptr[CBuffer] c_mask
2035+
20272036
cdef CMemoryPool* cpool = maybe_unbox_memory_pool(pool)
20282037

20292038
_offsets = asarray(offsets, type='int64')
20302039
_values = asarray(values)
20312040

2041+
c_mask = c_mask_inverted_from_obj(mask, pool)
2042+
20322043
if type is not None:
20332044
with nogil:
20342045
out = GetResultValue(
20352046
CLargeListArray.FromArraysAndType(
2036-
type.sp_type, _offsets.ap[0], _values.ap[0], cpool))
2047+
type.sp_type, _offsets.ap[0], _values.ap[0], cpool, c_mask))
20372048
else:
20382049
with nogil:
20392050
out = GetResultValue(
20402051
CLargeListArray.FromArrays(
2041-
_offsets.ap[0], _values.ap[0], cpool))
2052+
_offsets.ap[0], _values.ap[0], cpool, c_mask))
20422053
cdef Array result = pyarrow_wrap_array(out)
20432054
result.validate()
20442055
return result
@@ -2699,17 +2710,7 @@ cdef class StructArray(Array):
26992710
if names is not None and fields is not None:
27002711
raise ValueError('Must pass either names or fields, not both')
27012712

2702-
if mask is None:
2703-
c_mask = shared_ptr[CBuffer]()
2704-
elif isinstance(mask, Array):
2705-
if mask.type.id != Type_BOOL:
2706-
raise ValueError('Mask must be a pyarrow.Array of type bool')
2707-
if mask.null_count != 0:
2708-
raise ValueError('Mask must not contain nulls')
2709-
inverted_mask = _pc().invert(mask, memory_pool=memory_pool)
2710-
c_mask = pyarrow_unwrap_buffer(inverted_mask.buffers()[1])
2711-
else:
2712-
raise ValueError('Mask must be a pyarrow.Array of type bool')
2713+
c_mask = c_mask_inverted_from_obj(mask, memory_pool)
27132714

27142715
arrays = [asarray(x) for x in arrays]
27152716
for arr in arrays:
@@ -2859,6 +2860,25 @@ cdef dict _array_classes = {
28592860
}
28602861

28612862

2863+
cdef inline shared_ptr[CBuffer] c_mask_inverted_from_obj(object mask, MemoryPool pool) except *:
2864+
"""
2865+
Convert mask array obj to c_mask while also inverting to signify 1 for valid and 0 for null
2866+
"""
2867+
cdef shared_ptr[CBuffer] c_mask
2868+
if mask is None:
2869+
c_mask = shared_ptr[CBuffer]()
2870+
elif isinstance(mask, Array):
2871+
if mask.type.id != Type_BOOL:
2872+
raise TypeError('Mask must be a pyarrow.Array of type boolean')
2873+
if mask.null_count != 0:
2874+
raise ValueError('Mask must not contain nulls')
2875+
inverted_mask = _pc().invert(mask, memory_pool=pool)
2876+
c_mask = pyarrow_unwrap_buffer(inverted_mask.buffers()[1])
2877+
else:
2878+
raise TypeError('Mask must be a pyarrow.Array of type boolean')
2879+
return c_mask
2880+
2881+
28622882
cdef object get_array_class_from_type(
28632883
const shared_ptr[CDataType]& sp_data_type):
28642884
cdef CDataType* data_type = sp_data_type.get()

‎python/pyarrow/includes/libarrow.pxd‎

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -606,12 +606,20 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
606606
cdef cppclass CListArray" arrow::ListArray"(CArray):
607607
@staticmethod
608608
CResult[shared_ptr[CArray]] FromArrays(
609-
const CArray& offsets, const CArray& values, CMemoryPool* pool)
609+
const CArray& offsets,
610+
const CArray& values,
611+
CMemoryPool* pool,
612+
shared_ptr[CBuffer] null_bitmap,
613+
)
610614

611615
@staticmethod
612616
CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"(
613-
shared_ptr[CDataType], const CArray& offsets, const CArray& values,
614-
CMemoryPool* pool)
617+
shared_ptr[CDataType],
618+
const CArray& offsets,
619+
const CArray& values,
620+
CMemoryPool* pool,
621+
shared_ptr[CBuffer] null_bitmap,
622+
)
615623

616624
const int32_t* raw_value_offsets()
617625
int32_t value_offset(int i)
@@ -623,12 +631,20 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
623631
cdef cppclass CLargeListArray" arrow::LargeListArray"(CArray):
624632
@staticmethod
625633
CResult[shared_ptr[CArray]] FromArrays(
626-
const CArray& offsets, const CArray& values, CMemoryPool* pool)
634+
const CArray& offsets,
635+
const CArray& values,
636+
CMemoryPool* pool,
637+
shared_ptr[CBuffer] null_bitmap
638+
)
627639

628640
@staticmethod
629641
CResult[shared_ptr[CArray]] FromArraysAndType" FromArrays"(
630-
shared_ptr[CDataType], const CArray& offsets, const CArray& values,
631-
CMemoryPool* pool)
642+
shared_ptr[CDataType],
643+
const CArray& offsets,
644+
const CArray& values,
645+
CMemoryPool* pool,
646+
shared_ptr[CBuffer] null_bitmap
647+
)
632648

633649
int64_t value_offset(int i)
634650
int64_t value_length(int i)

0 commit comments

Comments
 (0)