ngoldbaum commented on code in PR #50951:
URL: https://github.com/apache/arrow/pull/50951#discussion_r3856824290
##########
python/pyarrow/src/arrow/python/numpy_to_arrow.cc:
##########
@@ -697,8 +706,76 @@ Status AppendUTF32(const char* data, int64_t itemsize, int
byteorder, T* builder
} // namespace
+template <typename T>
+Status NumPyConverter::VisitStringDType(T* builder) {
+ if (length_ == 0) {
+ return Status::OK();
+ }
+
+ const char* data = PyArray_BYTES(arr_);
+ auto* allocator =
+
NpyString_acquire_allocator(reinterpret_cast<PyArray_StringDTypeObject*>(dtype_));
+ std::unique_ptr<npy_string_allocator, decltype(&NpyString_release_allocator)>
+ allocator_guard(allocator, &NpyString_release_allocator);
+ const bool has_mask = mask_ != nullptr;
+ Ndarray1DIndexer<uint8_t> mask_values;
+ if (has_mask) {
+ mask_values = Ndarray1DIndexer<uint8_t>(mask_);
+ }
+
+ constexpr int64_t kBatchSize = 4096;
+ // NpyString_load returns borrowed views that remain valid while the
allocator is
+ // locked.
+ const int64_t batch_capacity = std::min(length_, kBatchSize);
+ std::vector<std::string_view> values(batch_capacity);
+ std::vector<uint8_t> valid(batch_capacity);
+ npy_static_string value{};
+
+ for (int64_t offset = 0; offset < length_; offset += kBatchSize) {
Review Comment:
```suggestion
const char* data = PyArray_BYTES(arr_);
const bool has_mask = mask_ != nullptr;
Ndarray1DIndexer<uint8_t> mask_values;
if (has_mask) {
mask_values = Ndarray1DIndexer<uint8_t>(mask_);
}
constexpr int64_t kBatchSize = 4096;
// NpyString_load returns borrowed views that remain valid while the
allocator is
// locked.
const int64_t batch_capacity = std::min(length_, kBatchSize);
std::vector<std::string_view> values(batch_capacity);
std::vector<uint8_t> valid(batch_capacity);
npy_static_string value{};
// Hold the allocator lock for the whole conversion to ensure a consistent
snapshot
auto* allocator =
NpyString_acquire_allocator(reinterpret_cast<PyArray_StringDTypeObject*>(dtype_));
std::unique_ptr<npy_string_allocator,
decltype(&NpyString_release_allocator)>
allocator_guard(allocator, &NpyString_release_allocator);
for (int64_t offset = 0; offset < length_; offset += kBatchSize) {
```
This way you hold the allocator lock for the minimum amount of code, I also
added a comment explaining that there is a mutex here so a reader knows to
watch out for deadlock risks.
##########
python/pyarrow/src/arrow/python/numpy_to_arrow.cc:
##########
@@ -697,8 +706,76 @@ Status AppendUTF32(const char* data, int64_t itemsize, int
byteorder, T* builder
} // namespace
+template <typename T>
+Status NumPyConverter::VisitStringDType(T* builder) {
+ if (length_ == 0) {
+ return Status::OK();
+ }
+
+ const char* data = PyArray_BYTES(arr_);
+ auto* allocator =
+
NpyString_acquire_allocator(reinterpret_cast<PyArray_StringDTypeObject*>(dtype_));
+ std::unique_ptr<npy_string_allocator, decltype(&NpyString_release_allocator)>
+ allocator_guard(allocator, &NpyString_release_allocator);
+ const bool has_mask = mask_ != nullptr;
+ Ndarray1DIndexer<uint8_t> mask_values;
+ if (has_mask) {
+ mask_values = Ndarray1DIndexer<uint8_t>(mask_);
+ }
+
+ constexpr int64_t kBatchSize = 4096;
+ // NpyString_load returns borrowed views that remain valid while the
allocator is
+ // locked.
+ const int64_t batch_capacity = std::min(length_, kBatchSize);
+ std::vector<std::string_view> values(batch_capacity);
+ std::vector<uint8_t> valid(batch_capacity);
+ npy_static_string value{};
+
+ for (int64_t offset = 0; offset < length_; offset += kBatchSize) {
+ const int64_t batch_length = std::min(kBatchSize, length_ - offset);
+ int64_t null_count = 0;
+
+ for (int64_t i = 0; i < batch_length; ++i) {
+ const char* item = data;
+ data += stride_;
+ if (has_mask && mask_values[offset + i]) {
+ values[i] = {};
+ valid[i] = false;
+ ++null_count;
+ continue;
+ }
+
+ const auto* packed = reinterpret_cast<const
npy_packed_static_string*>(item);
+ const int is_null = NpyString_load(allocator, packed, &value);
+ if (is_null == -1) {
+ return Status::Invalid("Failed to load NumPy StringDType value");
+ }
+ valid[i] = !is_null;
+ if (is_null) {
+ values[i] = {};
+ ++null_count;
+ } else {
+ values[i] = std::string_view(value.buf, value.size);
+ }
+ }
+
+ if (null_count == batch_length) {
+ RETURN_NOT_OK(builder->AppendNulls(batch_length));
+ } else {
+ RETURN_NOT_OK(builder->AppendValues(
+ std::span<const std::string_view>(values.data(), batch_length),
+ null_count == 0 ? NULLPTR : valid.data()));
+ }
+ }
+ return Status::OK();
+}
+
template <typename T>
Status NumPyConverter::VisitString(T* builder) {
+ if (dtype_->type_num == NPY_VSTRING) {
+ return VisitStringDType(builder);
Review Comment:
```suggestion
// Acquires a lock so this must not be moved inside the gil_lock section
below
return VisitStringDType(builder);
```
##########
python/pyarrow/src/arrow/python/numpy_to_arrow.cc:
##########
@@ -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(
+ "NumPy StringDType can only be converted to Arrow string types");
Review Comment:
Might be nice to see the type that was passed in:
```suggestion
"NumPy StringDType can only be converted to Arrow string types, got",
type_->ToString());
```
##########
cpp/src/arrow/array/builder_binary.h:
##########
@@ -23,7 +23,7 @@
#include <cstring>
#include <limits>
#include <memory>
-#include <numeric>
Review Comment:
This is a public header so this shouldn't be deleted. Deleting this also
caused churn in this PR in other compilation units, which newly add `#include
<numeric>` to work around this. The same could happen in user code including
this header.
##########
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);
Review Comment:
My AI model thinks that doing this has an unintended, untested side effect:
`pa.array([np.array([...], dtype="T")])` now infers as a list of strings and
converts through the per-element sequence fallback in `python_to_arrow.cc`:
https://github.com/apache/arrow/blob/54ea6d5babb6a1841dbe50e1cb84ae6321cd64cb/python/pyarrow/src/arrow/python/python_to_arrow.cc#L945-L947
That's fine, it just needs a test to cover it.
##########
python/pyarrow/tests/test_array.py:
##########
@@ -2925,6 +2925,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()
+
+ for sliced in (arr[:0], arr[::2], arr[::-1]):
+ arrow_arr = pa.array(sliced, type=string_type)
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == sliced.tolist()
+
+
[email protected]
[email protected]('string_type', [
+ None,
+ pa.large_string(),
+ pa.string_view(),
+])
[email protected]('na_object', [None, "__placeholder__", float("nan")])
+def test_array_from_numpy_string_dtype_nulls_and_mask(
+ numpy_string_dtype, string_type, na_object):
+ arr = np.array(["some", na_object, "strings"],
+ dtype=numpy_string_dtype(na_object=na_object))
+
+ arrow_arr = pa.array(arr, type=string_type)
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == ["some", None, "strings"]
Review Comment:
Why is the output missing sentinel `None` for all input `na_object` choices,
in particular for strings?
Note that this also disagrees with NumPy:
```
>>> np.array(["hello", "__placeholder__", "world"], dtype="T").tolist()
['hello', '__placeholder__', 'world']
```
Is there a reason why you can't more faithfully translate NumPy's string
missing data semantics to arrow's?
##########
python/pyarrow/tests/test_array.py:
##########
@@ -2925,6 +2925,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()
+
+ for sliced in (arr[:0], arr[::2], arr[::-1]):
+ arrow_arr = pa.array(sliced, type=string_type)
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == sliced.tolist()
+
+
[email protected]
[email protected]('string_type', [
+ None,
+ pa.large_string(),
+ pa.string_view(),
+])
[email protected]('na_object', [None, "__placeholder__", float("nan")])
+def test_array_from_numpy_string_dtype_nulls_and_mask(
+ numpy_string_dtype, string_type, na_object):
+ arr = np.array(["some", na_object, "strings"],
+ dtype=numpy_string_dtype(na_object=na_object))
+
+ arrow_arr = pa.array(arr, type=string_type)
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == ["some", None, "strings"]
+
+ mask = np.array([False, False, True])
+ arrow_arr = pa.array(arr, mask=mask, type=string_type)
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == ["some", None, None]
+
+
[email protected]
[email protected]('string_type', [pa.large_string(), pa.string_view()])
Review Comment:
```suggestion
@pytest.mark.parametrize('string_type', [None, pa.large_string(),
pa.string_view()])
```
Let's include the default `string()` type too.
##########
python/pyarrow/tests/test_array.py:
##########
@@ -2925,6 +2925,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()
+
+ for sliced in (arr[:0], arr[::2], arr[::-1]):
+ arrow_arr = pa.array(sliced, type=string_type)
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == sliced.tolist()
+
+
[email protected]
[email protected]('string_type', [
+ None,
+ pa.large_string(),
+ pa.string_view(),
+])
[email protected]('na_object', [None, "__placeholder__", float("nan")])
+def test_array_from_numpy_string_dtype_nulls_and_mask(
+ numpy_string_dtype, string_type, na_object):
+ arr = np.array(["some", na_object, "strings"],
+ dtype=numpy_string_dtype(na_object=na_object))
+
+ arrow_arr = pa.array(arr, type=string_type)
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == ["some", None, "strings"]
+
+ mask = np.array([False, False, True])
+ arrow_arr = pa.array(arr, mask=mask, type=string_type)
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == ["some", None, None]
+
+
[email protected]
[email protected]('string_type', [pa.large_string(), pa.string_view()])
+def test_array_from_numpy_string_dtype_batches(
+ numpy_string_dtype, string_type):
+ values = [None] * 4096 + [f"value-{i}" for i in range(904)]
+ arr = np.array(values, dtype=numpy_string_dtype(na_object=None))
+
+ arrow_arr = pa.array(arr, type=string_type)
+
+ arrow_arr.validate(full=True)
+ assert arrow_arr.to_pylist() == values
+
Review Comment:
The first test improves coverage for your changes to `ChunkedBinaryBuilder`.
The second test covers the untested case I referred to in my comment above
where a list of ndarrays gains new inference behavior.
```suggestion
@pytest.mark.numpy
def test_array_from_numpy_string_dtype_chunking(numpy_string_dtype):
# Three 6 MiB values in one batch must split across the 16 MiB
# per-chunk limit of the string() path.
values = ["x" * (6 * 1024 * 1024)] * 3
arr = np.array(values, dtype=numpy_string_dtype())
result = pa.array(arr, type=pa.string())
assert isinstance(result, pa.ChunkedArray)
assert result.num_chunks == 2
result.validate(full=True)
assert result.to_pylist() == values
@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
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
