rok commented on code in PR #48710:
URL: https://github.com/apache/arrow/pull/48710#discussion_r4444597212
##########
cpp/src/arrow/array/util.cc:
##########
@@ -927,6 +928,62 @@ Result<std::shared_ptr<Array>>
MakeEmptyArray(std::shared_ptr<DataType> type,
return builder->Finish();
}
+Result<std::shared_ptr<Array>> Arange(int64_t start, int64_t stop, int64_t
step,
+ MemoryPool* pool) {
+ if (step == 0) {
+ return Status::Invalid("Arange step cannot be zero");
+ }
+
+ auto error = [&]() {
+ return Status::Invalid("Arange: range [", start, ", ", stop, ") with step
", step,
+ " would produce an array that is too large");
+ };
+
+ int64_t size;
+ if (step > 0 && stop > start) {
+ // size = (stop - start + step - 1) / step
+ int64_t range;
+ if (internal::SubtractWithOverflow(stop, start, &range)) {
+ return error();
+ }
+ int64_t range_adjusted;
+ if (internal::AddWithOverflow(range, step, &range_adjusted) ||
+ internal::SubtractWithOverflow(range_adjusted, 1, &range_adjusted)) {
+ return error();
+ }
+ size = range_adjusted / step;
+
+ } else if (step < 0 && stop < start) {
+ if (step == INT64_MIN) {
+ return Status::Invalid("Arange step cannot be INT64_MIN");
+ }
+ // size = (start - stop - step - 1) / -step
+ int64_t range;
+ if (internal::SubtractWithOverflow(start, stop, &range)) {
+ return error();
+ }
+ int64_t range_adjusted;
+ if (internal::SubtractWithOverflow(range, step, &range_adjusted) ||
+ internal::SubtractWithOverflow(range_adjusted, 1, &range_adjusted)) {
+ return error();
+ }
+ size = range_adjusted / -step;
+
+ } else {
+ return MakeEmptyArray(int64(), pool);
+ }
+
+ ARROW_ASSIGN_OR_RAISE(auto data_buffer, AllocateBuffer(size *
sizeof(int64_t), pool));
Review Comment:
This looks like it could overflow indeed. We could handle it here with a
`MultiplyWithOverflow` or perhaps we can calculate that a buffer will overflow
given `start`, `step` and `stop` parameters and throw at the top of this
function?
--
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]
