Reranko05 commented on code in PR #48710:
URL: https://github.com/apache/arrow/pull/48710#discussion_r4444438284
##########
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:
@rok I was reviewing this PR and had a small doubt. Should we also check for
potential overflow when calculating the allocation size? The number of elements
might fit in an int64, but multiplying it by the element size could still
overflow. Are we handling this case somewhere else?
--
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]
