diff --git a/docs/source/python/numpy.rst b/docs/source/python/numpy.rst index 07a6aa803f8..10a70159c8e 100644 --- a/docs/source/python/numpy.rst +++ b/docs/source/python/numpy.rst @@ -51,7 +51,28 @@ factory function. ] Converting from NumPy supports a wide range of input dtypes, including -structured dtypes or strings. +structured dtypes and both fixed-width (``S`` and ``U``) and variable-width +(:class:`numpy.dtypes.StringDType`) strings. + +A ``StringDType`` array converts to :func:`~pyarrow.string` unless +:func:`~pyarrow.large_string` or :func:`~pyarrow.string_view` is requested with +``type``. Missing entries become nulls: + +.. code-block:: python + + >>> dtype = np.dtypes.StringDType(na_object=np.nan) + >>> arr = pa.array(np.array(["some", np.nan, "strings"], dtype=dtype)) + >>> arr + + [ + "some", + null, + "strings" + ] + +When the ``na_object`` is a string, NumPy treats missing entries as that string +in every operation, and so does the conversion. Pass ``mask`` to mark values as +null explicitly. Arrow to NumPy -------------- diff --git a/python/pyarrow/src/arrow/python/numpy_convert.cc b/python/pyarrow/src/arrow/python/numpy_convert.cc index 6e59835286d..882e9e47e5b 100644 --- a/python/pyarrow/src/arrow/python/numpy_convert.cc +++ b/python/pyarrow/src/arrow/python/numpy_convert.cc @@ -151,6 +151,7 @@ Result> NumPyDtypeToArrow(PyArray_Descr* descr) { TO_ARROW_TYPE_CASE(FLOAT64, float64); TO_ARROW_TYPE_CASE(STRING, binary); TO_ARROW_TYPE_CASE(UNICODE, utf8); + TO_ARROW_TYPE_CASE(VSTRING, utf8); case NPY_DATETIME: { auto date_dtype = reinterpret_cast(PyDataType_C_METADATA(descr)); diff --git a/python/pyarrow/src/arrow/python/numpy_to_arrow.cc b/python/pyarrow/src/arrow/python/numpy_to_arrow.cc index 5647e895d0f..d4e859d1c70 100644 --- a/python/pyarrow/src/arrow/python/numpy_to_arrow.cc +++ b/python/pyarrow/src/arrow/python/numpy_to_arrow.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include #include @@ -295,6 +296,9 @@ class NumPyConverter { template Status VisitString(T* builder); + template + Status VisitStringDType(T* builder); + Status TypeNotImplemented(std::string type_name) { return Status::NotImplemented("NumPyConverter doesn't implement <", type_name, "> conversion. "); @@ -342,6 +346,11 @@ Status NumPyConverter::Convert() { return Status::Invalid("Must pass data type for non-object arrays"); } + if (dtype_->type_num == NPY_VSTRING && !is_string_or_string_view(type_->id())) { + return Status::TypeError("Expected an Arrow string type for NumPy StringDType, got ", + type_->ToString()); + } + // Visit the type to perform conversion return VisitTypeInline(*type_, this); } @@ -697,8 +706,65 @@ Status AppendUTF32(const char* data, int64_t itemsize, int byteorder, T* builder } // namespace +namespace { + +std::string_view ToStringView(const npy_static_string& value) { + return value.buf == nullptr ? std::string_view() + : std::string_view(value.buf, value.size); +} + +} // namespace + +template +Status NumPyConverter::VisitStringDType(T* builder) { + auto* descr = reinterpret_cast(dtype_); + // Use the na_object itself when na_object is a string + const bool null_is_missing = descr->na_object != nullptr && !descr->has_string_na; + const std::string_view null_string = ToStringView(descr->default_string); + + const char* data = PyArray_BYTES(arr_); + Ndarray1DIndexer mask_values; + if (mask_ != nullptr) { + mask_values = Ndarray1DIndexer(mask_); + } + + // Acquiring the allocator lock, so do not acquire the GIL or lock other + // mutexes below or risk deadlocks + auto* allocator = NpyString_acquire_allocator(descr); + std::unique_ptr + allocator_guard(allocator, &NpyString_release_allocator); + + npy_static_string value = {0, nullptr}; + for (int64_t i = 0; i < length_; ++i, data += stride_) { + if (mask_ != nullptr && mask_values[i]) { + RETURN_NOT_OK(builder->AppendNull()); + continue; + } + const auto* packed = reinterpret_cast(data); + const int is_null = NpyString_load(allocator, packed, &value); + if (is_null == -1) { + return Status::Invalid("Failed to load NumPy StringDType value"); + } + if (is_null) { + if (null_is_missing) { + RETURN_NOT_OK(builder->AppendNull()); + } else { + RETURN_NOT_OK(builder->Append(null_string)); + } + continue; + } + RETURN_NOT_OK(builder->Append(ToStringView(value))); + } + return Status::OK(); +} + template Status NumPyConverter::VisitString(T* builder) { + if (dtype_->type_num == NPY_VSTRING) { + // Acquires a lock, so must stay ahead of the gil_lock below + return VisitStringDType(builder); + } + auto data = reinterpret_cast(PyArray_DATA(arr_)); char numpy_byteorder = dtype_->byteorder; diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index a1e3616c9ce..062248d1269 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -2925,6 +2925,85 @@ def test_array_from_numpy_unicode(string_type): assert arrow_arr.equals(expected) +@pytest.fixture +def numpy_string_dtype(): + dtypes = pytest.importorskip("numpy.dtypes") + return dtypes.StringDType + + +@pytest.mark.numpy +@pytest.mark.parametrize('string_type', [ + None, + pa.string(), + pa.large_string(), + pa.string_view()]) +def test_array_from_numpy_string_dtype(numpy_string_dtype, string_type): + values = [ + "short", + "a" * 100, + "b" * 300, + "árvíztűrő tükörfúrógép 🥐 你好", + "🥐" * 200, + "", + ] + arr = np.array(values, dtype=numpy_string_dtype()) + + arrow_arr = pa.array(arr, type=string_type) + + arrow_arr.validate(full=True) + assert arrow_arr.type == (string_type or pa.string()) + assert arrow_arr.to_pylist() == arr.tolist() + + strided = np.array(list(itertools.chain.from_iterable( + zip(values, itertools.repeat("skip")))), + dtype=numpy_string_dtype())[::2] + arrow_arr = pa.array(strided, type=string_type) + arrow_arr.validate(full=True) + assert arrow_arr.to_pylist() == values + + +@pytest.mark.numpy +@pytest.mark.parametrize('na_object, expected', [ + (None, None), + (float("nan"), None), + ("__placeholder__", "__placeholder__"), +]) +def test_array_from_numpy_string_dtype_na_object( + numpy_string_dtype, na_object, expected): + arr = np.array(["some", na_object, "strings"], + dtype=numpy_string_dtype(na_object=na_object)) + + arrow_arr = pa.array(arr) + arrow_arr.validate(full=True) + assert arrow_arr.to_pylist() == ["some", expected, "strings"] + + mask = np.array([False, False, True]) + arrow_arr = pa.array(arr, mask=mask) + arrow_arr.validate(full=True) + assert arrow_arr.to_pylist() == ["some", expected, None] + + +@pytest.mark.numpy +def test_array_from_numpy_string_dtype_rejects_non_string_type( + numpy_string_dtype): + arr = np.array(["some", "strings"], dtype=numpy_string_dtype()) + + msg = "Expected an Arrow string type.*got binary" + with pytest.raises(TypeError, match=msg): + pa.array(arr, type=pa.binary()) + + +@pytest.mark.numpy +def test_array_from_list_of_numpy_string_dtype_arrays(numpy_string_dtype): + values = [["a", "bb"], ["ccc"]] + arrays = [np.array(v, dtype=numpy_string_dtype()) for v in values] + + result = pa.array(arrays) + + assert result.type == pa.list_(pa.string()) + assert result.to_pylist() == values + + @pytest.mark.numpy def test_array_string_from_non_string(): # ARROW-5682 - when converting to string raise on non string-like dtype