BUG: Improve NA handling in to/from wkb/t (#1891)

geopandas/_vectorized.py

@@ -7,6 +7,7 @@ Uses PyGEOS if available/set, otherwise loops through Shapely geometries.
 import warnings
 
 import numpy as np
+import pandas as pd
 
 import shapely.geometry
 import shapely.geos
@@ -44,9 +45,9 @@ else:
     type_mapping, geometry_type_ids, geometry_type_values = None, None, None
 
 
-def _isna(value):
+def isna(value):
     """
-    Check if scalar value is NA-like (None or np.nan).
+    Check if scalar value is NA-like (None, np.nan or pd.NA).
 
     Custom version that only works for scalars (returning True or False),
     as `pd.isna` also works for array-like input returning a boolean array.
@@ -55,6 +56,8 @@ def _isna(value):
         return True
     elif isinstance(value, float) and np.isnan(value):
         return True
+    elif compat.PANDAS_GE_10 and value is pd.NA:
+        return True
     else:
         return False
 
@@ -127,7 +130,7 @@ def from_shapely(data):
                 out.append(_shapely_to_pygeos(geom))
             else:
                 out.append(geom)
-        elif _isna(geom):
+        elif isna(geom):
             out.append(None)
         else:
             raise TypeError("Input must be valid geometry objects: {0}".format(geom))
@@ -165,7 +168,7 @@ def from_wkb(data):
     out = []
 
     for geom in data:
-        if geom is not None and len(geom):
+        if not isna(geom) and len(geom):
             geom = shapely.wkb.loads(geom)
         else:
             geom = None
@@ -200,7 +203,7 @@ def from_wkt(data):
     out = []
 
     for geom in data:
-        if geom is not None and len(geom):
+        if not isna(geom) and len(geom):
             if isinstance(geom, bytes):
                 geom = geom.decode("utf-8")
             geom = shapely.wkt.loads(geom)
@@ -941,7 +944,7 @@ def transform(data, func):
         result = np.empty(n, dtype=object)
         for i in range(n):
             geom = data[i]
-            if _isna(geom):
+            if isna(geom):
                 result[i] = geom
             else:
                 result[i] = transform(func, geom)

geopandas/array.py

@@ -56,23 +56,6 @@ class GeometryDtype(ExtensionDtype):
 register_extension_dtype(GeometryDtype)
 
 
-def _isna(value):
-    """
-    Check if scalar value is NA-like (None, np.nan or pd.NA).
-
-    Custom version that only works for scalars (returning True or False),
-    as `pd.isna` also works for array-like input returning a boolean array.
-    """
-    if value is None:
-        return True
-    elif isinstance(value, float) and np.isnan(value):
-        return True
-    elif compat.PANDAS_GE_10 and value is pd.NA:
-        return True
-    else:
-        return False
-
-
 def _check_crs(left, right, allow_none=False):
     """
     Check if the projection of both arrays is the same.
@@ -398,8 +381,8 @@ class GeometryArray(ExtensionArray):
             if isinstance(key, numbers.Integral):
                 raise ValueError("cannot set a single element with an array")
             self.data[key] = value.data
-        elif isinstance(value, BaseGeometry) or _isna(value):
-            if _isna(value):
+        elif isinstance(value, BaseGeometry) or vectorized.isna(value):
+            if vectorized.isna(value):
                 # internally only use None as missing value indicator
                 # but accept others
                 value = None
@@ -1005,7 +988,7 @@ class GeometryArray(ExtensionArray):
 
         if mask.any():
             # fill with value
-            if _isna(value):
+            if vectorized.isna(value):
                 value = None
             elif not isinstance(value, BaseGeometry):
                 raise NotImplementedError(
@@ -1326,7 +1309,7 @@ class GeometryArray(ExtensionArray):
         """
         Return for `item in self`.
         """
-        if _isna(item):
+        if vectorized.isna(item):
             if (
                 item is self.dtype.na_value
                 or isinstance(item, self.dtype.type)

geopandas/tests/test_array.py

@@ -139,15 +139,16 @@ def test_from_wkb():
     assert all(v.equals(t) for v, t in zip(res, points_no_missing))
 
     # missing values
-    # TODO(pygeos) does not support empty strings
-    if compat.USE_PYGEOS:
-        L_wkb.extend([None])
-    else:
-        L_wkb.extend([b"", None])
-    res = from_wkb(L_wkb)
-    assert res[-1] is None
+    # TODO(pygeos) does not support empty strings, np.nan, or pd.NA
+    missing_values = [None]
     if not compat.USE_PYGEOS:
-        assert res[-2] is None
+        missing_values.extend([b"", np.nan])
+
+        if compat.PANDAS_GE_10:
+            missing_values.append(pd.NA)
+
+    res = from_wkb(missing_values)
+    np.testing.assert_array_equal(res, np.full(len(missing_values), None))
 
     # single MultiPolygon
     multi_poly = shapely.geometry.MultiPolygon(
@@ -202,15 +203,16 @@ def test_from_wkt(string_type):
     assert all(v.almost_equals(t) for v, t in zip(res, points_no_missing))
 
     # missing values
-    # TODO(pygeos) does not support empty strings
-    if compat.USE_PYGEOS:
-        L_wkt.extend([None])
-    else:
-        L_wkt.extend([f(""), None])
-    res = from_wkt(L_wkt)
-    assert res[-1] is None
+    # TODO(pygeos) does not support empty strings, np.nan, or pd.NA
+    missing_values = [None]
     if not compat.USE_PYGEOS:
-        assert res[-2] is None
+        missing_values.extend([f(""), np.nan])
+
+        if compat.PANDAS_GE_10:
+            missing_values.append(pd.NA)
+
+    res = from_wkb(missing_values)
+    np.testing.assert_array_equal(res, np.full(len(missing_values), None))
 
     # single MultiPolygon
     multi_poly = shapely.geometry.MultiPolygon(