GH-51652: [Python] Add NumPy StringDType to Arrow conversion

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
 --------------

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));

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;

python/pyarrow/tests/test_array.py

@@ -2983,6 +2983,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

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: