Skip to content

Commit b524c85

Browse files
committed
Fix zero-sized storeChunk for Span API in Python
1 parent a258128 commit b524c85

File tree

1 file changed

+27
-9
lines changed

1 file changed

+27
-9
lines changed

src/binding/python/RecordComponent.cpp

Lines changed: 27 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -529,7 +529,7 @@ struct PythonDynamicMemoryView
529529
, m_datatype(determineDatatype<T>())
530530
{}
531531

532-
[[nodiscard]] pybind11::memoryview currentView() const;
532+
[[nodiscard]] pybind11::object currentView() const;
533533

534534
std::shared_ptr<void> m_dynamicView;
535535
ShapeContainer m_arrayShape;
@@ -542,30 +542,48 @@ namespace
542542
struct GetCurrentView
543543
{
544544
template <typename T>
545-
static pybind11::memoryview call(PythonDynamicMemoryView const &dynamicView)
545+
static pybind11::object call(PythonDynamicMemoryView const &dynamicView)
546546
{
547547
auto span =
548548
static_cast<DynamicMemoryView<T> *>(dynamicView.m_dynamicView.get())
549549
->currentBuffer();
550-
return py::memoryview::from_buffer(
551-
span.data(),
552-
dynamicView.m_arrayShape,
553-
dynamicView.m_strides,
554-
/* readonly = */ false);
550+
if (!span.data())
551+
{
552+
/*
553+
* Fallback for zero-sized store_chunk calls.
554+
* py::memoryview cannot be used since it checks for nullpointers,
555+
* even when the extent is zero.
556+
* This may sound like an esoteric use case at first, but may happen
557+
* in parallel usage when the chunk distribution ends up assigning
558+
* zero-sized chunks to some rank.
559+
*/
560+
return py::array(
561+
dtype_to_numpy(dynamicView.m_datatype),
562+
dynamicView.m_arrayShape,
563+
dynamicView.m_strides);
564+
}
565+
else
566+
{
567+
return py::memoryview::from_buffer(
568+
span.data(),
569+
dynamicView.m_arrayShape,
570+
dynamicView.m_strides,
571+
/* readonly = */ false);
572+
}
555573
}
556574

557575
static constexpr char const *errorMsg = "DynamicMemoryView";
558576
};
559577

560578
template <>
561-
pybind11::memoryview
579+
pybind11::object
562580
GetCurrentView::call<std::string>(PythonDynamicMemoryView const &)
563581
{
564582
throw std::runtime_error("[DynamicMemoryView] Only PODs allowed.");
565583
}
566584
} // namespace
567585

568-
pybind11::memoryview PythonDynamicMemoryView::currentView() const
586+
pybind11::object PythonDynamicMemoryView::currentView() const
569587
{
570588
return switchNonVectorType<GetCurrentView>(m_datatype, *this);
571589
}

0 commit comments

Comments
 (0)