This is an automated email from the ASF dual-hosted git repository.
jorisvandenbossche pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 17be9c7caa9 GH-51652: [Python] Add NumPy StringDType to Arrow
conversion (#51157)
17be9c7caa9 is described below
commit 17be9c7caa928cb3ec5a1dfbd2c2037e78720637
Author: Nathan Goldbaum <[email protected]>
AuthorDate: Thu Oct 1 16:24:28 2026 +0300
GH-51652: [Python] Add NumPy StringDType to Arrow conversion (#51157)
### Rationale for this change
NumPy 2.0 added `StringDType`, a variable-width UTF-8 string dtype.
`pa.array` currently rejects it with `Unsupported numpy type 2056`. Now that
pyarrow requires NumPy 2.0 the conversion can use the public `NpyString_*` C
API.
This supersedes #50951. The first commit is @ alippai's first commit from
that PR, unchanged. His later commits added batching and new bulk builder APIs
in Arrow C++; those can come later in their own PR with benchmarks, as
discussed there. Nothing outside the NumPy conversion code changes here.
### What changes are included in this PR?
- `NumPyDtypeToArrow` maps `StringDType` to `string`, so it is also the
inferred type for `pa.array`, `pa.infer_type`, `pa.from_numpy_dtype` and lists
of such arrays. `large_string` and `string_view` can be requested with `type`.
Any other type raises a `TypeError` naming the requested type.
- The conversion holds the dtype's allocator lock for the whole array and
never holds the GIL while taking it.
- A null entry in a StringDType array becomes an Arrow null unless the
dtype's `na_object` is a `str`. NumPy reports such entries as that string from
`__getitem__`, in ufuncs and comparisons, and in casts, and there is no way to
tell them apart from a regular entry holding the same string, so the conversion
writes the string too. NaN-like sentinels such as `np.nan` and `pd.NA`, `None`,
and arbitrary objects all become arrow nulls.
- A short section in `docs/source/python/numpy.rst` describing the support
and the semantics for converting missing data.
If left unfixed, https://github.com/apache/arrow/issues/51156 also becomes
possible to trigger from Python by creating a StringDType array with a very
large entry.
### Are these changes tested?
Yes. Tests cover the three target types, strings crossing NumPy's short,
medium and long storage thresholds, non-ASCII input, strided input, `mask`,
`None`, NaN and string sentinels, the rejected target type, and list inference.
I also ran a threaded test that converts an array while other threads hit
NumPy's null-comparison error path on the same array to check the lock ordering.
### Are there any user-facing changes?
Yes. `pa.array` and the inference functions accept `StringDType` arrays.
AI disclosure: I used an AI model to work on the follow-up commits on top
of @ alippai's first commit and for code review.
* GitHub Issue: #42018
* GitHub Issue: #51652
Lead-authored-by: Nathan Goldbaum <[email protected]>
Co-authored-by: Joris Van den Bossche <[email protected]>
Co-authored-by: Adam Lippai <[email protected]>
Signed-off-by: Joris Van den Bossche <[email protected]>
---
docs/source/python/numpy.rst | 23 ++++++-
python/pyarrow/src/arrow/python/numpy_convert.cc | 1 +
python/pyarrow/src/arrow/python/numpy_to_arrow.cc | 66 +++++++++++++++++++
python/pyarrow/tests/test_array.py | 79 +++++++++++++++++++++++
python/pyarrow/tests/test_schema.py | 1 +
5 files changed, 169 insertions(+), 1 deletion(-)
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
+ <pyarrow.lib.StringArray object at ...>
+ [
+ "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<std::shared_ptr<DataType>>
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<PyArray_DatetimeDTypeMetaData*>(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 <limits>
#include <memory>
#include <string>
+#include <string_view>
#include <utility>
#include <vector>
@@ -295,6 +296,9 @@ class NumPyConverter {
template <typename T>
Status VisitString(T* builder);
+ template <typename T>
+ 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 <typename T>
+Status NumPyConverter::VisitStringDType(T* builder) {
+ auto* descr = reinterpret_cast<PyArray_StringDTypeObject*>(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<uint8_t> mask_values;
+ if (mask_ != nullptr) {
+ mask_values = Ndarray1DIndexer<uint8_t>(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<npy_string_allocator, decltype(&NpyString_release_allocator)>
+ 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<const
npy_packed_static_string*>(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 <typename T>
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<const uint8_t*>(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 87040c0800f..38a131f5d95 100644
--- a/python/pyarrow/tests/test_array.py
+++ b/python/pyarrow/tests/test_array.py
@@ -2983,6 +2983,85 @@ def test_array_from_numpy_unicode(string_type):
assert arrow_arr.equals(expected)
[email protected]
+def numpy_string_dtype():
+ dtypes = pytest.importorskip("numpy.dtypes")
+ return dtypes.StringDType
+
+
[email protected]
[email protected]('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
+
+
[email protected]
[email protected]('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]
+
+
[email protected]
+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())
+
+
[email protected]
+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
diff --git a/python/pyarrow/tests/test_schema.py
b/python/pyarrow/tests/test_schema.py
index 93519fc0173..b35b5350a81 100644
--- a/python/pyarrow/tests/test_schema.py
+++ b/python/pyarrow/tests/test_schema.py
@@ -209,6 +209,7 @@ def test_from_numpy_dtype():
(np.dtype('timedelta64[ms]'), pa.duration('ms')),
(np.dtype('timedelta64[us]'), pa.duration('us')),
(np.dtype('timedelta64[ns]'), pa.duration('ns')),
+ (np.dtypes.StringDType(), pa.string()),
]
for dt, pt in cases:
