diff --git a/python/pyarrow/src/arrow/python/arrow_to_pandas.cc b/python/pyarrow/src/arrow/python/arrow_to_pandas.cc index eb60d9d54f8..d5ca924ebf0 100644 --- a/python/pyarrow/src/arrow/python/arrow_to_pandas.cc +++ b/python/pyarrow/src/arrow/python/arrow_to_pandas.cc @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include @@ -747,6 +748,64 @@ Status ConvertStruct(PandasOptions options, const ChunkedArray& data, return Status::OK(); } +// Whether decoding `arr` into its string/binary value type would overflow the +// 32-bit offsets of the dense array (GH-50842), i.e. whether Take would reject +// it. The length of the dense data is computed from the dictionary offsets. +bool DecodingWouldOverflow(const DictionaryArray& arr) { + // As in the Take kernel for binary-like arrays + constexpr int64_t kOffsetLimit = std::numeric_limits::max() - 1; + const auto& dictionary = checked_cast(*arr.dictionary()); + const int64_t dictionary_length = dictionary.length(); + // Cheap bounds first: no overflow is possible if every row referred to the + // whole dictionary, or if no value is longer than kOffsetLimit / length + const int64_t length = arr.length(); + const int64_t dictionary_bytes = + dictionary.value_offset(dictionary_length) - dictionary.value_offset(0); + if (dictionary_bytes <= 0 || length <= kOffsetLimit / dictionary_bytes) { + return false; + } + if (dictionary_length <= length) { + const int64_t max_value_length = kOffsetLimit / length; + int64_t i = 0; + while (i < dictionary_length && dictionary.value_length(i) <= max_value_length) { + ++i; + } + if (i == dictionary_length) { + return false; + } + } + const bool has_nulls = arr.null_count() > 0; + int64_t total_length = 0; + for (int64_t i = 0; i < length; ++i) { + if (has_nulls && arr.IsNull(i)) { + continue; + } + const int64_t index = arr.GetValueIndex(i); + if (index < 0 || index >= dictionary_length) { + // Let the decoding report the invalid index + return false; + } + total_length += dictionary.value_length(index); + if (total_length > kOffsetLimit) { + return true; + } + } + return false; +} + +bool DecodingWouldOverflow(const ChunkedArray& arr) { + const auto& value_type = checked_cast(*arr.type()).value_type(); + if (!is_binary_like(value_type->id())) { + return false; + } + for (int c = 0; c < arr.num_chunks(); c++) { + if (DecodingWouldOverflow(checked_cast(*arr.chunk(c)))) { + return true; + } + } + return false; +} + Status DecodeDictionaries(MemoryPool* pool, const std::shared_ptr& dense_type, ArrayVector* arrays) { compute::ExecContext ctx(pool); @@ -1380,6 +1439,40 @@ struct ObjectWriterVisitor { return ConvertStruct(options, data, out_values); } + // Dictionaries whose decoding would overflow the 32-bit offsets of the dense + // array (GH-50842) are not decoded. Instead, the dictionary values are + // converted once and referenced for each index. + Status Visit(const DictionaryType& type) { + for (int c = 0; c < data.num_chunks(); c++) { + const auto& arr = checked_cast(*data.chunk(c)); + const int64_t dictionary_length = arr.dictionary()->length(); + std::vector values(dictionary_length, nullptr); + ChunkedArray dictionary(arr.dictionary()); + ObjectWriterVisitor values_visitor{options, dictionary, values.data()}; + Status st = VisitTypeInline(*arr.dictionary()->type(), &values_visitor); + for (int64_t i = 0; st.ok() && i < arr.length(); ++i) { + if (arr.IsNull(i)) { + Py_INCREF(Py_None); + *out_values = Py_None; + } else { + const int64_t index = arr.GetValueIndex(i); + if (index < 0 || index >= dictionary_length) { + st = Status::IndexError("Index ", index, " out of bounds"); + break; + } + Py_INCREF(values[index]); + *out_values = values[index]; + } + ++out_values; + } + for (PyObject* value : values) { + Py_XDECREF(value); + } + RETURN_NOT_OK(st); + } + return Status::OK(); + } + Status Visit(const ExtensionType& type) { if (!IsUuidExtension(type)) { return Status::NotImplemented("No implemented conversion to object dtype: ", @@ -2275,9 +2368,16 @@ static Status GetPandasWriterType(const ChunkedArray& data, const PandasOptions& } *output_type = PandasWriter::OBJECT; } break; - case Type::DICTIONARY: - *output_type = PandasWriter::CATEGORICAL; - break; + case Type::DICTIONARY: { + // A string/binary dictionary that is not decoded because the dense array + // would overflow is converted to objects instead, see + // ObjectWriterVisitor::Visit(const DictionaryType&) + const auto& value_type = + checked_cast(*data.type()).value_type(); + *output_type = options.decode_dictionaries && is_binary_like(value_type->id()) + ? PandasWriter::OBJECT + : PandasWriter::CATEGORICAL; + } break; case Type::EXTENSION: // UUID has a native object conversion to uuid.UUID. Other extension // types continue through the pandas ExtensionArray protocol. @@ -2595,7 +2695,11 @@ Status ConvertArrayToPandas(const PandasOptions& options, std::shared_ptr Status ConvertChunkedArrayToPandas(const PandasOptions& options, std::shared_ptr arr, PyObject* py_ref, PyObject** out) { - if (options.decode_dictionaries && arr->type()->id() == Type::DICTIONARY) { + // Decoding a string/binary dictionary can overflow the 32-bit offsets of + // the dense array (GH-50842): such arrays are converted to objects without + // decoding instead, see ObjectWriterVisitor::Visit(const DictionaryType&) + if (options.decode_dictionaries && arr->type()->id() == Type::DICTIONARY && + !DecodingWouldOverflow(*arr)) { // XXX we should return an error as below if options.zero_copy_only // is true, but that would break compatibility with existing tests. const auto& dense_type = diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 74ff837e16e..c5933fe8949 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -1082,6 +1082,38 @@ def test_dictionary_to_numpy(): ) +@pytest.mark.numpy +def test_dictionary_to_numpy_without_decoding(): + # GH-50842: a dictionary whose decoded form would overflow the 32-bit + # offsets of its value type (16385 * 2**17 > INT32_MAX bytes here) is + # converted to objects without decoding + value = "a" * 2**17 + arr = pa.DictionaryArray.from_arrays( + np.zeros(16385, dtype=np.int16), pa.array([value])) + result = np.asarray(arr) + assert len(result) == len(arr) + assert result[0] == value + assert result[-1] == value + + # With a null index (row 2) and a null dictionary value (row 1) + indices = np.zeros(16387, dtype=np.int16) + indices[1] = 1 + arr = pa.DictionaryArray.from_arrays( + pa.array(indices, mask=np.arange(16387) == 2), pa.array([value, None])) + result = arr.to_numpy(zero_copy_only=False) + assert result[0] == value + assert result[1] is None + assert result[2] is None + assert result[-1] == value + + # Binary values + arr = pa.DictionaryArray.from_arrays( + np.zeros(16385, dtype=np.int16), pa.array([value.encode()])) + result = np.asarray(arr) + assert result[0] == value.encode() + assert result[-1] == value.encode() + + @pytest.mark.numpy def test_dictionary_from_boxed_arrays(): indices = np.repeat([0, 1, 2], 2) diff --git a/python/pyarrow/tests/test_pandas.py b/python/pyarrow/tests/test_pandas.py index 9421b345269..5f64cb1bd3e 100644 --- a/python/pyarrow/tests/test_pandas.py +++ b/python/pyarrow/tests/test_pandas.py @@ -2518,6 +2518,24 @@ def test_list_of_dictionary(self): expected[2] = None tm.assert_series_equal(arr.to_pandas(), expected) + def test_list_of_dictionary_without_decoding(self): + # GH-50842: the dictionary child is converted to objects without + # being decoded into a dense array, which would overflow its 32-bit + # offsets (16385 * 2**17 > INT32_MAX bytes here) + value = "a" * 2**17 + child = pa.DictionaryArray.from_arrays( + np.zeros(16385, dtype=np.int16), pa.array([value])) + arr = pa.ListArray.from_arrays([0, len(child)], child) + result = arr.to_pandas() + assert len(result[0]) == len(child) + assert result[0][0] == value + + # Also with several chunks + result = pa.chunked_array([arr, arr]).to_pandas() + assert len(result) == 2 + assert len(result[1]) == len(child) + assert result[1][-1] == value + @pytest.mark.large_memory def test_auto_chunking_on_list_overflow(self): # ARROW-9976