diff --git a/cpp/src/arrow/array/util.cc b/cpp/src/arrow/array/util.cc index 03d8c32c4e3..629fc1c6d57 100644 --- a/cpp/src/arrow/array/util.cc +++ b/cpp/src/arrow/array/util.cc @@ -927,6 +927,32 @@ Result> MakeEmptyArray(std::shared_ptr type, return builder->Finish(); } +Result> Arange(int64_t start, int64_t stop, int64_t step, + MemoryPool* pool) { + if (step == 0) { + return Status::Invalid("arange: step cannot be zero"); + } + + int64_t size; + if (step > 0 && stop > start) { + size = (stop - start + step - 1) / step; + } else if (step < 0 && stop < start) { + size = (start - stop - step - 1) / (-step); + } else { + return MakeEmptyArray(int64(), pool); + } + + ARROW_ASSIGN_OR_RAISE(auto data_buffer, AllocateBuffer(size * sizeof(int64_t), pool)); + + auto values = reinterpret_cast(data_buffer->mutable_data()); + for (int64_t i = 0; i < size; ++i) { + values[i] = start + i * step; + } + + auto data = ArrayData::Make(int64(), size, {nullptr, std::move(data_buffer)}, 0); + return MakeArray(data); +} + namespace internal { std::vector RechunkArraysConsistently( diff --git a/cpp/src/arrow/array/util.h b/cpp/src/arrow/array/util.h index fd8e75ddb86..2e9a992c252 100644 --- a/cpp/src/arrow/array/util.h +++ b/cpp/src/arrow/array/util.h @@ -69,6 +69,21 @@ ARROW_EXPORT Result> MakeEmptyArray(std::shared_ptr type, MemoryPool* pool = default_memory_pool()); +/// \brief Create an array of evenly spaced integer values +/// +/// Generates a sequence of integers from start (inclusive) to stop (exclusive) +/// with the given step, similar to Python's range() or NumPy's arange(). +/// The resulting array will have type Int64. +/// +/// \param[in] start First value in the sequence +/// \param[in] stop End value (exclusive) +/// \param[in] step Spacing between values (must not be zero) +/// \param[in] pool the memory pool to allocate memory from +/// \return the resulting Int64Array +ARROW_EXPORT +Result> Arange(int64_t start, int64_t stop, int64_t step, + MemoryPool* pool = default_memory_pool()); + /// @} namespace internal { diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index d80f21333e7..4229c64c9d3 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -375,8 +375,7 @@ set(PYARROW_CPP_SRCS ${PYARROW_CPP_SOURCE_DIR}/python_test.cc ${PYARROW_CPP_SOURCE_DIR}/python_to_arrow.cc ${PYARROW_CPP_SOURCE_DIR}/pyarrow.cc - ${PYARROW_CPP_SOURCE_DIR}/udf.cc - ${PYARROW_CPP_SOURCE_DIR}/util.cc) + ${PYARROW_CPP_SOURCE_DIR}/udf.cc) set_source_files_properties(${PYARROW_CPP_SOURCE_DIR}/numpy_init.cc PROPERTIES SKIP_UNITY_BUILD_INCLUSION ON) diff --git a/python/pyarrow/includes/libarrow.pxd b/python/pyarrow/includes/libarrow.pxd index c03bf20026e..539ef94375d 100644 --- a/python/pyarrow/includes/libarrow.pxd +++ b/python/pyarrow/includes/libarrow.pxd @@ -287,6 +287,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil: CResult[shared_ptr[CArray]] MakeArrayFromScalar( const CScalar& scalar, int64_t length, CMemoryPool* pool) + CResult[shared_ptr[CArray]] Arange(int64_t start, int64_t stop, + int64_t step, CMemoryPool* pool) + CStatus DebugPrint(const CArray& arr, int indent) cdef cppclass CFixedWidthType" arrow::FixedWidthType"(CDataType): diff --git a/python/pyarrow/includes/libarrow_python.pxd b/python/pyarrow/includes/libarrow_python.pxd index 4724c52ccb5..613a63a1d30 100644 --- a/python/pyarrow/includes/libarrow_python.pxd +++ b/python/pyarrow/includes/libarrow_python.pxd @@ -73,9 +73,6 @@ cdef extern from "arrow/python/api.h" namespace "arrow::py" nogil: object obj, object mask, const PyConversionOptions& options, CMemoryPool* pool) - CResult[shared_ptr[CArray]] Arange(int64_t start, int64_t stop, - int64_t step, CMemoryPool* pool) - CResult[shared_ptr[CDataType]] NumPyDtypeToArrow(object dtype) CStatus NdarrayToArrow(CMemoryPool* pool, object ao, object mo, diff --git a/python/pyarrow/src/arrow/python/api.h b/python/pyarrow/src/arrow/python/api.h index 2af0963a9c0..e66bf49dfec 100644 --- a/python/pyarrow/src/arrow/python/api.h +++ b/python/pyarrow/src/arrow/python/api.h @@ -26,4 +26,3 @@ #include "arrow/python/numpy_convert.h" #include "arrow/python/numpy_to_arrow.h" #include "arrow/python/python_to_arrow.h" -#include "arrow/python/util.h" diff --git a/python/pyarrow/src/arrow/python/util.cc b/python/pyarrow/src/arrow/python/util.cc deleted file mode 100644 index cffe1eb956b..00000000000 --- a/python/pyarrow/src/arrow/python/util.cc +++ /dev/null @@ -1,50 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -#include "arrow/python/util.h" - -#include "arrow/array.h" -#include "arrow/python/common.h" - -namespace arrow ::py { - -Result> Arange(int64_t start, int64_t stop, int64_t step, - MemoryPool* pool) { - int64_t size; - if (step == 0) { - return Status::Invalid("Step must not be zero"); - } - if (step > 0 && stop > start) { - // Ceiling division for positive step - size = (stop - start + step - 1) / step; - } else if (step < 0 && stop < start) { - // Ceiling division for negative step - size = (start - stop - step - 1) / (-step); - } else { - return MakeEmptyArray(int64()); - } - std::shared_ptr data_buffer; - ARROW_ASSIGN_OR_RAISE(data_buffer, AllocateBuffer(size * sizeof(int64_t), pool)); - auto values = reinterpret_cast(data_buffer->mutable_data()); - for (int64_t i = 0; i < size; ++i) { - values[i] = start + i * step; - } - auto data = ArrayData::Make(int64(), size, {nullptr, data_buffer}, 0); - return MakeArray(data); -} - -} // namespace arrow::py diff --git a/python/pyarrow/src/arrow/python/util.h b/python/pyarrow/src/arrow/python/util.h deleted file mode 100644 index ff2ffcaea9c..00000000000 --- a/python/pyarrow/src/arrow/python/util.h +++ /dev/null @@ -1,40 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -#pragma once - -#include "arrow/python/common.h" -#include "arrow/python/visibility.h" - -namespace arrow::py { - -/// \brief Create an array of evenly spaced values within a given interval. -/// This function is similar to Python's `range` function. -/// The resulting array will contain values starting from `start` up to but not -/// including `stop`, with a step size of `step`. If `step` is zero, the function -/// will return an error. -/// The resulting array will have a data type of `int64`. -/// \param[in] start initial value of the sequence. -/// \param[in] stop final value of the sequence (exclusive). -/// \param[in] step step size between consecutive values. -/// \param[in] pool Memory pool for any memory allocations. -/// \return Result Array -ARROW_PYTHON_EXPORT -Result> Arange(int64_t start, int64_t stop, int64_t step, - MemoryPool* pool); - -} // namespace arrow::py