Finish first pass as fixes

rapidsai · Aug 7, 2024 · 5f9be85 · 5f9be85
1 parent 1825283
commit 5f9be85
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 16 deletions.
diff --git a/python/cuspatial/cuspatial/core/_column/geometa.py b/python/cuspatial/cuspatial/core/_column/geometa.py
@@ -7,6 +7,7 @@
 from typing import Literal, Union
 
 import cudf
+import cudf.core.column
 
 
 # This causes arrow to encode NONE as =255, which I'll accept now
@@ -37,12 +38,22 @@ def __init__(
         ],
     ):
         if isinstance(meta, dict):
-            self.input_types = cudf.Series._from_column(
-                meta["input_types"]
-            ).astype("int8")
-            self.union_offsets = cudf.Series._from_column(
-                meta["union_offsets"]
-            ).astype("int32")
+            meta_it = meta["input_types"]
+            if isinstance(meta_it, cudf.core.column.ColumnBase):
+                self.input_types = cudf.Series._from_column(meta_it).astype(
+                    "int8"
+                )
+            else:
+                # Could be Series from GeoSeries.__getitem__
+                self.input_types = cudf.Series(meta_it, dtype="int8")
+            meta_uo = meta["union_offsets"]
+            if isinstance(meta_uo, cudf.core.column.ColumnBase):
+                self.union_offsets = cudf.Series._from_column(meta_uo).astype(
+                    "int8"
+                )
+            else:
+                # Could be Series from GeoSeries.__getitem__
+                self.union_offsets = cudf.Series(meta_uo, dtype="int8")
         else:
             self.input_types = cudf.Series(meta.input_types, dtype="int8")
             self.union_offsets = cudf.Series(meta.union_offsets, dtype="int32")

diff --git a/python/cuspatial/cuspatial/core/geoseries.py b/python/cuspatial/cuspatial/core/geoseries.py
@@ -87,24 +87,25 @@ def __init__(
             data = gpGeoSeries(data)
         # Create column
         if isinstance(data, GeoSeries):
-            gser = type(self)._from_column(data._column, name=data.name)
+            gser = type(self)._from_column(
+                data._column, index=data.index, name=data.name
+            )
         elif isinstance(data, gpGeoSeries):
             from cuspatial.io.geopandas_reader import GeoPandasReader
 
             adapter = GeoPandasReader(data)
             pandas_meta = GeoMeta(adapter.get_geopandas_meta())
             geocolumn = GeoColumn(adapter._get_geotuple(), pandas_meta)
-            gser = type(self)._from_column(geocolumn, name=data.name)
+            gser = type(self)._from_column(
+                geocolumn, index=cudf.Index(data.index), name=data.name
+            )
         else:
             raise TypeError(
                 f"Incompatible object passed to GeoSeries ctor {type(data)}"
             )
         # Condition index
-        if isinstance(data, (gpGeoSeries, GeoSeries)):
-            if index is None:
-                gser.index = data.index
-        if index is None:
-            gser.index = cudf.RangeIndex(len(gser))
+        if index is not None:
+            gser.index = cudf.Index(index)
         if name is not None:
             gser.name = name
         super().__init__(gser, dtype=dtype, nan_as_null=nan_as_null)
@@ -637,7 +638,7 @@ def to_arrow(self):
 
         return pa.UnionArray.from_dense(
             self._column._meta.input_types.to_arrow(),
-            self._column._meta.union_offsets.to_arrow(),
+            self._column._meta.union_offsets.astype("int32").to_arrow(),
             [
                 arrow_points,
                 arrow_mpoints,

diff --git a/python/cuspatial/cuspatial/core/trajectory.py b/python/cuspatial/cuspatial/core/trajectory.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2019-2022, NVIDIA CORPORATION.
+# Copyright (c) 2019-2024, NVIDIA CORPORATION.
 
 import numpy as np
 
@@ -71,7 +71,7 @@ def derive_trajectories(object_ids, points: GeoSeries, timestamps):
     objects, traj_offsets = cpp_derive_trajectories(
         object_ids, xs, ys, timestamps
     )
-    return DataFrame._from_data(*objects), Series(data=traj_offsets)
+    return DataFrame._from_data(*objects), Series._from_column(traj_offsets)
 
 
 def trajectory_bounding_boxes(num_trajectories, object_ids, points: GeoSeries):