Updates cudf usage for compatibility with cudf>=23.10.1

rapidsai · Nov 16, 2023 · 43e20e4 · 43e20e4
1 parent bbc8fed
commit 43e20e4
Showing 1 changed file with 6 additions and 18 deletions.
diff --git a/python/cugraph/cugraph/structure/property_graph.py b/python/cugraph/cugraph/structure/property_graph.py
@@ -800,15 +800,9 @@ def add_vertex_data(
             tmp_df.index = tmp_df.index.rename(self.vertex_col_name)
 
         # FIXME: handle case of a type_name column already being in tmp_df
-        if self.__series_type is cudf.Series:
-            # cudf does not yet support initialization with a scalar
-            tmp_df[TCN] = cudf.Series(
-                cudf.Series([type_name], dtype=cat_dtype).repeat(len(tmp_df)),
-                index=tmp_df.index,
-            )
-        else:
-            # pandas is oddly slow if dtype is passed to the constructor here
-            tmp_df[TCN] = pd.Series(type_name, index=tmp_df.index).astype(cat_dtype)
+        tmp_df[TCN] = self.__series_type(type_name, index=tmp_df.index).astype(
+            cat_dtype
+        )
 
         if property_columns:
             # all columns
@@ -1207,15 +1201,9 @@ def add_edge_data(
         tmp_df[self.src_col_name] = tmp_df[vertex_col_names[0]]
         tmp_df[self.dst_col_name] = tmp_df[vertex_col_names[1]]
 
-        if self.__series_type is cudf.Series:
-            # cudf does not yet support initialization with a scalar
-            tmp_df[TCN] = cudf.Series(
-                cudf.Series([type_name], dtype=cat_dtype).repeat(len(tmp_df)),
-                index=tmp_df.index,
-            )
-        else:
-            # pandas is oddly slow if dtype is passed to the constructor here
-            tmp_df[TCN] = pd.Series(type_name, index=tmp_df.index).astype(cat_dtype)
+        tmp_df[TCN] = self.__series_type(type_name, index=tmp_df.index).astype(
+            cat_dtype
+        )
 
         # Add unique edge IDs to the new rows. This is just a count for each
         # row starting from the last edge ID value, with initial edge ID 0.