Speed up geometry encoding

Closes #534
xarray-contrib · Sep 10, 2024 · d8d3de8 · d8d3de8
1 parent acc1eb0
commit d8d3de8
Showing 1 changed file with 24 additions and 21 deletions.
diff --git a/cf_xarray/geometry.py b/cf_xarray/geometry.py
@@ -460,10 +460,11 @@ def shapely_to_cf(
         )
 
     # Get all types to call the appropriate translation function.
-    types = {
-        geom.item().geom_type if isinstance(geom, xr.DataArray) else geom.geom_type
-        for geom in geometries
-    }
+    geometries = (
+        geometries.data if isinstance(geometries, xr.DataArray) else geometries
+    )  #
+    # types = {geom.geom_type for geom in geometries}
+    types = {geometries[0].geom_type}
 
     grid_mapping_varname = None
     if (
@@ -841,7 +842,7 @@ def polygons_to_cf(
         node_count = part_node_count
     elif len(offsets) >= 2:
         indices = np.take(offsets[0], offsets[1])
-        interior_ring = np.isin(offsets[0], indices, invert=True)[:-1].astype(int)
+        interior_ring = np.isin(offsets[0], indices, invert=True)[:-1].view(np.int8)
 
         if len(offsets) == 3:
             indices = np.take(indices, offsets[2])
@@ -852,29 +853,31 @@ def polygons_to_cf(
     crdX = geom_coords[:, 0]
     crdY = geom_coords[:, 1]
 
+    data_vars = {names.node_count: (dim, node_count)}
+
+    # Special case when we have no MultiPolygons and no holes
+    if len(part_node_count) != len(node_count):
+        data_vars[names.part_node_count] = (names.part_dim, part_node_count)
+        names.geometry_container_attrs["part_node_count"] = names.part_node_count
+
+    # Special case when we have no holes
+    if (interior_ring != 0).any():
+        data_vars[names.interior_ring] = (names.part_dim, interior_ring)
+        names.geometry_container_attrs["interior_ring"] = names.interior_ring
+
+    data_vars[names.container_name] = (
+        (),
+        np.nan,
+        {"geometry_type": "polygon", **names.geometry_container_attrs},
+    )
     ds = xr.Dataset(
-        data_vars={
-            names.node_count: xr.DataArray(node_count, dims=(dim,)),
-            names.container_name: xr.DataArray(
-                data=np.nan,
-                attrs={"geometry_type": "polygon", **names.geometry_container_attrs},
-            ),
-        },
+        data_vars=data_vars,
         coords=names.coords(x=x, y=y, crdX=crdX, crdY=crdY, dim=dim),
     )
 
     if coord is not None:
         ds = ds.assign_coords({dim: coord})
 
-    # Special case when we have no MultiPolygons and no holes
-    if len(part_node_count) != len(node_count):
-        ds[names.part_node_count] = xr.DataArray(part_node_count, dims=names.part_dim)
-        ds[names.container_name].attrs["part_node_count"] = names.part_node_count
-
-    # Special case when we have no holes
-    if (interior_ring != 0).any():
-        ds[names.interior_ring] = xr.DataArray(interior_ring, dims=names.part_dim)
-        ds[names.container_name].attrs["interior_ring"] = names.interior_ring
     return ds