| 23 | namespace arrow ::py { |
| 24 | |
| 25 | Result<std::shared_ptr<Array>> Arange(int64_t start, int64_t stop, int64_t step, |
| 26 | MemoryPool* pool) { |
| 27 | int64_t size; |
| 28 | if (step == 0) { |
| 29 | return Status::Invalid("Step must not be zero"); |
| 30 | } |
| 31 | if (step > 0 && stop > start) { |
| 32 | // Ceiling division for positive step |
| 33 | size = (stop - start + step - 1) / step; |
| 34 | } else if (step < 0 && stop < start) { |
| 35 | // Ceiling division for negative step |
| 36 | size = (start - stop - step - 1) / (-step); |
| 37 | } else { |
| 38 | return MakeEmptyArray(int64()); |
| 39 | } |
| 40 | std::shared_ptr<Buffer> data_buffer; |
| 41 | ARROW_ASSIGN_OR_RAISE(data_buffer, AllocateBuffer(size * sizeof(int64_t), pool)); |
| 42 | auto values = reinterpret_cast<int64_t*>(data_buffer->mutable_data()); |
| 43 | for (int64_t i = 0; i < size; ++i) { |
| 44 | values[i] = start + i * step; |
| 45 | } |
| 46 | auto data = ArrayData::Make(int64(), size, {nullptr, data_buffer}, 0); |
| 47 | return MakeArray(data); |
| 48 | } |
| 49 | |
| 50 | } // namespace arrow::py |
nothing calls this directly
no test coverage detected