Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 22 additions & 1 deletion docs/source/python/numpy.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
--------------
Expand Down
1 change: 1 addition & 0 deletions python/pyarrow/src/arrow/python/numpy_convert.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand Down
66 changes: 66 additions & 0 deletions python/pyarrow/src/arrow/python/numpy_to_arrow.cc
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
#include <limits>
#include <memory>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

Expand Down Expand Up @@ -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. ");
Expand Down Expand Up @@ -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);
}
Expand Down Expand Up @@ -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;
Expand Down
79 changes: 79 additions & 0 deletions python/pyarrow/tests/test_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading