Reranko05 commented on code in PR #48710:
URL: https://github.com/apache/arrow/pull/48710#discussion_r4444658145
##########
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:
Agreed, `MultiplyWithOverflow` might be a good fit here. Do you think this
would be within the scope of this PR, or would it be better to open a separate
issue for it?
--
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]
