Skip to content

Commit 17be9c7

Browse files
ngoldbaumjorisvandenbosschealippai
authored
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, #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 <nathan.goldbaum@gmail.com> Co-authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com> Co-authored-by: Adam Lippai <adam@rigo.sk> Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
1 parent 6f161d6 commit 17be9c7

5 files changed

Lines changed: 169 additions & 1 deletion

File tree

‎docs/source/python/numpy.rst‎

Lines changed: 22 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,28 @@ factory function.
5151
]
5252
5353
Converting from NumPy supports a wide range of input dtypes, including
54-
structured dtypes or strings.
54+
structured dtypes and both fixed-width (``S`` and ``U``) and variable-width
55+
(:class:`numpy.dtypes.StringDType`) strings.
56+
57+
A ``StringDType`` array converts to :func:`~pyarrow.string` unless
58+
:func:`~pyarrow.large_string` or :func:`~pyarrow.string_view` is requested with
59+
``type``. Missing entries become nulls:
60+
61+
.. code-block:: python
62+
63+
>>> dtype = np.dtypes.StringDType(na_object=np.nan)
64+
>>> arr = pa.array(np.array(["some", np.nan, "strings"], dtype=dtype))
65+
>>> arr
66+
<pyarrow.lib.StringArray object at ...>
67+
[
68+
"some",
69+
null,
70+
"strings"
71+
]
72+
73+
When the ``na_object`` is a string, NumPy treats missing entries as that string
74+
in every operation, and so does the conversion. Pass ``mask`` to mark values as
75+
null explicitly.
5576

5677
Arrow to NumPy
5778
--------------

‎python/pyarrow/src/arrow/python/numpy_convert.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ Result<std::shared_ptr<DataType>> NumPyDtypeToArrow(PyArray_Descr* descr) {
151151
TO_ARROW_TYPE_CASE(FLOAT64, float64);
152152
TO_ARROW_TYPE_CASE(STRING, binary);
153153
TO_ARROW_TYPE_CASE(UNICODE, utf8);
154+
TO_ARROW_TYPE_CASE(VSTRING, utf8);
154155
case NPY_DATETIME: {
155156
auto date_dtype =
156157
reinterpret_cast<PyArray_DatetimeDTypeMetaData*>(PyDataType_C_METADATA(descr));

‎python/pyarrow/src/arrow/python/numpy_to_arrow.cc‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
#include <limits>
2828
#include <memory>
2929
#include <string>
30+
#include <string_view>
3031
#include <utility>
3132
#include <vector>
3233

@@ -295,6 +296,9 @@ class NumPyConverter {
295296
template <typename T>
296297
Status VisitString(T* builder);
297298

299+
template <typename T>
300+
Status VisitStringDType(T* builder);
301+
298302
Status TypeNotImplemented(std::string type_name) {
299303
return Status::NotImplemented("NumPyConverter doesn't implement <", type_name,
300304
"> conversion. ");
@@ -342,6 +346,11 @@ Status NumPyConverter::Convert() {
342346
return Status::Invalid("Must pass data type for non-object arrays");
343347
}
344348

349+
if (dtype_->type_num == NPY_VSTRING && !is_string_or_string_view(type_->id())) {
350+
return Status::TypeError("Expected an Arrow string type for NumPy StringDType, got ",
351+
type_->ToString());
352+
}
353+
345354
// Visit the type to perform conversion
346355
return VisitTypeInline(*type_, this);
347356
}
@@ -697,8 +706,65 @@ Status AppendUTF32(const char* data, int64_t itemsize, int byteorder, T* builder
697706

698707
} // namespace
699708

709+
namespace {
710+
711+
std::string_view ToStringView(const npy_static_string& value) {
712+
return value.buf == nullptr ? std::string_view()
713+
: std::string_view(value.buf, value.size);
714+
}
715+
716+
} // namespace
717+
718+
template <typename T>
719+
Status NumPyConverter::VisitStringDType(T* builder) {
720+
auto* descr = reinterpret_cast<PyArray_StringDTypeObject*>(dtype_);
721+
// Use the na_object itself when na_object is a string
722+
const bool null_is_missing = descr->na_object != nullptr && !descr->has_string_na;
723+
const std::string_view null_string = ToStringView(descr->default_string);
724+
725+
const char* data = PyArray_BYTES(arr_);
726+
Ndarray1DIndexer<uint8_t> mask_values;
727+
if (mask_ != nullptr) {
728+
mask_values = Ndarray1DIndexer<uint8_t>(mask_);
729+
}
730+
731+
// Acquiring the allocator lock, so do not acquire the GIL or lock other
732+
// mutexes below or risk deadlocks
733+
auto* allocator = NpyString_acquire_allocator(descr);
734+
std::unique_ptr<npy_string_allocator, decltype(&NpyString_release_allocator)>
735+
allocator_guard(allocator, &NpyString_release_allocator);
736+
737+
npy_static_string value = {0, nullptr};
738+
for (int64_t i = 0; i < length_; ++i, data += stride_) {
739+
if (mask_ != nullptr && mask_values[i]) {
740+
RETURN_NOT_OK(builder->AppendNull());
741+
continue;
742+
}
743+
const auto* packed = reinterpret_cast<const npy_packed_static_string*>(data);
744+
const int is_null = NpyString_load(allocator, packed, &value);
745+
if (is_null == -1) {
746+
return Status::Invalid("Failed to load NumPy StringDType value");
747+
}
748+
if (is_null) {
749+
if (null_is_missing) {
750+
RETURN_NOT_OK(builder->AppendNull());
751+
} else {
752+
RETURN_NOT_OK(builder->Append(null_string));
753+
}
754+
continue;
755+
}
756+
RETURN_NOT_OK(builder->Append(ToStringView(value)));
757+
}
758+
return Status::OK();
759+
}
760+
700761
template <typename T>
701762
Status NumPyConverter::VisitString(T* builder) {
763+
if (dtype_->type_num == NPY_VSTRING) {
764+
// Acquires a lock, so must stay ahead of the gil_lock below
765+
return VisitStringDType(builder);
766+
}
767+
702768
auto data = reinterpret_cast<const uint8_t*>(PyArray_DATA(arr_));
703769

704770
char numpy_byteorder = dtype_->byteorder;

‎python/pyarrow/tests/test_array.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2983,6 +2983,85 @@ def test_array_from_numpy_unicode(string_type):
29832983
assert arrow_arr.equals(expected)
29842984

29852985

2986+
@pytest.fixture
2987+
def numpy_string_dtype():
2988+
dtypes = pytest.importorskip("numpy.dtypes")
2989+
return dtypes.StringDType
2990+
2991+
2992+
@pytest.mark.numpy
2993+
@pytest.mark.parametrize('string_type', [
2994+
None,
2995+
pa.string(),
2996+
pa.large_string(),
2997+
pa.string_view()])
2998+
def test_array_from_numpy_string_dtype(numpy_string_dtype, string_type):
2999+
values = [
3000+
"short",
3001+
"a" * 100,
3002+
"b" * 300,
3003+
"árvíztűrő tükörfúrógép 🥐 你好",
3004+
"🥐" * 200,
3005+
"",
3006+
]
3007+
arr = np.array(values, dtype=numpy_string_dtype())
3008+
3009+
arrow_arr = pa.array(arr, type=string_type)
3010+
3011+
arrow_arr.validate(full=True)
3012+
assert arrow_arr.type == (string_type or pa.string())
3013+
assert arrow_arr.to_pylist() == arr.tolist()
3014+
3015+
strided = np.array(list(itertools.chain.from_iterable(
3016+
zip(values, itertools.repeat("skip")))),
3017+
dtype=numpy_string_dtype())[::2]
3018+
arrow_arr = pa.array(strided, type=string_type)
3019+
arrow_arr.validate(full=True)
3020+
assert arrow_arr.to_pylist() == values
3021+
3022+
3023+
@pytest.mark.numpy
3024+
@pytest.mark.parametrize('na_object, expected', [
3025+
(None, None),
3026+
(float("nan"), None),
3027+
("__placeholder__", "__placeholder__"),
3028+
])
3029+
def test_array_from_numpy_string_dtype_na_object(
3030+
numpy_string_dtype, na_object, expected):
3031+
arr = np.array(["some", na_object, "strings"],
3032+
dtype=numpy_string_dtype(na_object=na_object))
3033+
3034+
arrow_arr = pa.array(arr)
3035+
arrow_arr.validate(full=True)
3036+
assert arrow_arr.to_pylist() == ["some", expected, "strings"]
3037+
3038+
mask = np.array([False, False, True])
3039+
arrow_arr = pa.array(arr, mask=mask)
3040+
arrow_arr.validate(full=True)
3041+
assert arrow_arr.to_pylist() == ["some", expected, None]
3042+
3043+
3044+
@pytest.mark.numpy
3045+
def test_array_from_numpy_string_dtype_rejects_non_string_type(
3046+
numpy_string_dtype):
3047+
arr = np.array(["some", "strings"], dtype=numpy_string_dtype())
3048+
3049+
msg = "Expected an Arrow string type.*got binary"
3050+
with pytest.raises(TypeError, match=msg):
3051+
pa.array(arr, type=pa.binary())
3052+
3053+
3054+
@pytest.mark.numpy
3055+
def test_array_from_list_of_numpy_string_dtype_arrays(numpy_string_dtype):
3056+
values = [["a", "bb"], ["ccc"]]
3057+
arrays = [np.array(v, dtype=numpy_string_dtype()) for v in values]
3058+
3059+
result = pa.array(arrays)
3060+
3061+
assert result.type == pa.list_(pa.string())
3062+
assert result.to_pylist() == values
3063+
3064+
29863065
@pytest.mark.numpy
29873066
def test_array_string_from_non_string():
29883067
# ARROW-5682 - when converting to string raise on non string-like dtype

‎python/pyarrow/tests/test_schema.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,7 @@ def test_from_numpy_dtype():
209209
(np.dtype('timedelta64[ms]'), pa.duration('ms')),
210210
(np.dtype('timedelta64[us]'), pa.duration('us')),
211211
(np.dtype('timedelta64[ns]'), pa.duration('ns')),
212+
(np.dtypes.StringDType(), pa.string()),
212213
]
213214

214215
for dt, pt in cases:

0 commit comments

Comments
 (0)