Use as_column instead of arange for range like inputs

rapidsai · Jan 12, 2024 · a9818fe · a9818fe
1 parent 36549cd
commit a9818fe
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/python/cuspatial/cuspatial/core/_column/geocolumn.py b/python/cuspatial/cuspatial/core/_column/geocolumn.py
@@ -8,7 +8,7 @@
 import pyarrow as pa
 
 import cudf
-from cudf.core.column import ColumnBase, arange, as_column, build_list_column
+from cudf.core.column import ColumnBase, as_column, build_list_column
 
 from cuspatial.core._column.geometa import Feature_Enum, GeoMeta
 from cuspatial.utils.column_utils import empty_geometry_column
@@ -364,5 +364,5 @@ def _xy_as_variable_sized_list(xy: ColumnBase):
         raise ValueError("xy must have an even number of elements")
 
     num_points = len(xy) // 2
-    indices = arange(0, num_points * 2 + 1, 2, dtype="int32")
+    indices = as_column(range(0, num_points * 2 + 1, 2), dtype="int32")
     return build_list_column(indices=indices, elements=xy, size=num_points)
diff --git a/python/cuspatial/cuspatial/core/binops/distance_dispatch.py b/python/cuspatial/cuspatial/core/binops/distance_dispatch.py
@@ -1,5 +1,5 @@
 import cudf
-from cudf.core.column import arange, full
+from cudf.core.column import as_column, full
 
 from cuspatial._lib.distance import (
     pairwise_linestring_distance,
@@ -190,8 +190,8 @@ def __call__(self):
             float("nan"),
             dtype="float64",
         )
-        scatter_map = arange(
-            len(self._res_index), dtype="int32"
+        scatter_map = as_column(
+            range(len(self._res_index)), dtype="int32"
         ).apply_boolean_mask(self._non_null_mask)
 
         result[scatter_map] = dist

diff --git a/python/cuspatial/cuspatial/core/binops/intersection.py b/python/cuspatial/cuspatial/core/binops/intersection.py
@@ -3,7 +3,7 @@
 from typing import TYPE_CHECKING
 
 import cudf
-from cudf.core.column import arange, build_list_column
+from cudf.core.column import as_column, build_list_column
 
 from cuspatial._lib.intersection import (
     pairwise_linestring_intersection as c_pairwise_linestring_intersection,
@@ -93,7 +93,7 @@ def pairwise_linestring_intersection(
     ]
 
     linestring_column = build_list_column(
-        indices=arange(0, len(segments) + 1, dtype="int32"),
+        indices=as_column(range(0, len(segments) + 1), dtype="int32"),
         elements=segments,
         size=len(segments),
     )