Skip to content

Commit d3d6371

Browse files
authored
ARROW-14495: [Python] Fix DictionaryArray.from_buffers, should not crash (apache#13989)
Fix [ARROW-14495](https://issues.apache.org/jira/browse/ARROW-14495) Authored-by: Miles Granger <miles59923@gmail.com> Signed-off-by: Antoine Pitrou <antoine@python.org>
1 parent ec7e250 commit d3d6371

3 files changed

Lines changed: 57 additions & 0 deletions

File tree

‎python/pyarrow/array.pxi‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2478,6 +2478,53 @@ cdef class DictionaryArray(Array):
24782478

24792479
return self._indices
24802480

2481+
@staticmethod
2482+
def from_buffers(DataType type, int64_t length, buffers, Array dictionary,
2483+
int64_t null_count=-1, int64_t offset=0):
2484+
"""
2485+
Construct a DictionaryArray from buffers.
2486+
2487+
Parameters
2488+
----------
2489+
type : pyarrow.DataType
2490+
length : int
2491+
The number of values in the array.
2492+
buffers : List[Buffer]
2493+
The buffers backing the indices array.
2494+
dictionary : pyarrow.Array, ndarray or pandas.Series
2495+
The array of values referenced by the indices.
2496+
null_count : int, default -1
2497+
The number of null entries in the indices array. Negative value means that
2498+
the null count is not known.
2499+
offset : int, default 0
2500+
The array's logical offset (in values, not in bytes) from the
2501+
start of each buffer.
2502+
2503+
Returns
2504+
-------
2505+
dict_array : DictionaryArray
2506+
"""
2507+
cdef:
2508+
vector[shared_ptr[CBuffer]] c_buffers
2509+
shared_ptr[CDataType] c_type
2510+
shared_ptr[CArrayData] c_data
2511+
shared_ptr[CArray] c_result
2512+
2513+
for buf in buffers:
2514+
c_buffers.push_back(pyarrow_unwrap_buffer(buf))
2515+
2516+
c_type = pyarrow_unwrap_data_type(type)
2517+
2518+
with nogil:
2519+
c_data = CArrayData.Make(
2520+
c_type, length, c_buffers, null_count, offset)
2521+
c_data.get().dictionary = dictionary.sp_array.get().data()
2522+
c_result.reset(new CDictionaryArray(c_data))
2523+
2524+
cdef Array result = pyarrow_wrap_array(c_result)
2525+
result.validate()
2526+
return result
2527+
24812528
@staticmethod
24822529
def from_arrays(indices, dictionary, mask=None, bint ordered=False,
24832530
bint from_pandas=False, bint safe=True,

‎python/pyarrow/includes/libarrow.pxd‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -250,6 +250,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
250250
CDictionaryArray(const shared_ptr[CDataType]& type,
251251
const shared_ptr[CArray]& indices,
252252
const shared_ptr[CArray]& dictionary)
253+
CDictionaryArray(const shared_ptr[CArrayData]& data)
253254

254255
@staticmethod
255256
CResult[shared_ptr[CArray]] FromArrays(

‎python/pyarrow/tests/test_array.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -725,6 +725,15 @@ def test_struct_array_from_chunked():
725725
pa.StructArray.from_arrays([chunked_arr], ["foo"])
726726

727727

728+
@pytest.mark.parametrize("offset", (0, 1))
729+
def test_dictionary_from_buffers(offset):
730+
a = pa.array(["one", "two", "three", "two", "one"]).dictionary_encode()
731+
b = pa.DictionaryArray.from_buffers(a.type, len(a)-offset,
732+
a.indices.buffers(), a.dictionary,
733+
offset=offset)
734+
assert a[offset:] == b
735+
736+
728737
def test_dictionary_from_numpy():
729738
indices = np.repeat([0, 1, 2], 2)
730739
dictionary = np.array(['foo', 'bar', 'baz'], dtype=object)

0 commit comments

Comments
 (0)