Major speedups in mappers. (#427)

* Major speedups in mappers.

* fix test

Author: dcherian

cf_xarray/accessor.py

@@ -230,6 +230,7 @@ def _get_custom_criteria(
 
     if isinstance(obj, DataArray):
         obj = obj._to_temp_dataset()
+    variables = obj._variables
 
     if criteria is None:
         if not OPTIONS["custom_criteria"]:
@@ -243,8 +244,8 @@ def _get_custom_criteria(
     results: set = set()
     if key in criteria_map:
         for criterion, patterns in criteria_map[key].items():
-            for var in obj.variables:
-                if regex_match(patterns, obj[var].attrs.get(criterion, "")):
+            for var in variables:
+                if regex_match(patterns, variables[var].attrs.get(criterion, "")):
                     results.update((var,))
                 # also check name specifically since not in attributes
                 elif (
@@ -290,6 +291,8 @@ def _get_axis_coord(obj: DataArray | Dataset, key: str) -> list[str]:
             f"cf_xarray did not understand key {key!r}. Expected one of {valid_keys!r}"
         )
 
+    crds = obj.coords
+    crd_names = set(crds)
     search_in = set()
     attrs_or_encoding = ChainMap(obj.attrs, obj.encoding)
     coordinates = attrs_or_encoding.get("coordinates", None)
@@ -298,15 +301,16 @@ def _get_axis_coord(obj: DataArray | Dataset, key: str) -> list[str]:
     if coordinates:
         search_in.update(coordinates.split(" "))
     if not search_in:
-        search_in = set(obj.coords)
+        search_in = crd_names
 
     # maybe only do this for key in _AXIS_NAMES?
-    search_in.update(obj.indexes)
+    if obj._indexes:
+        search_in.update(obj._indexes)
 
-    search_in = search_in & set(obj.coords)
+    search_in = search_in & crd_names
     results: set = set()
     for coord in search_in:
-        var = obj.coords[coord]
+        var = crds[coord]
         if key in coordinate_criteria:
             for criterion, expected in coordinate_criteria[key].items():
                 if var.attrs.get(criterion, None) in expected:
@@ -345,9 +349,8 @@ def _get_measure(obj: DataArray | Dataset, key: str) -> list[str]:
         obj = obj._to_temp_dataset()
 
     results = set()
-    for var in obj.variables:
-        da = obj[var]
-        attrs_or_encoding = ChainMap(da.attrs, da.encoding)
+    for var in obj._variables.values():
+        attrs_or_encoding = ChainMap(var.attrs, var.encoding)
         if "cell_measures" in attrs_or_encoding:
             attr = attrs_or_encoding["cell_measures"]
             try:
@@ -381,9 +384,13 @@ def _get_bounds(obj: DataArray | Dataset, key: Hashable) -> list[Hashable]:
     List[str], Variable name(s) in parent xarray object that are bounds of `key`
     """
 
+    if isinstance(obj, DataArray):
+        obj = obj._to_temp_dataset()
+    variables = obj._variables
+
     results = set()
     for var in apply_mapper(_get_all, obj, key, error=False, default=[key]):
-        attrs_or_encoding = ChainMap(obj[var].attrs, obj[var].encoding)
+        attrs_or_encoding = ChainMap(variables[var].attrs, variables[var].encoding)
         if "bounds" in attrs_or_encoding:
             results |= {attrs_or_encoding["bounds"]}
 
@@ -410,17 +417,17 @@ def _get_grid_mapping_name(obj: DataArray | Dataset, key: str) -> list[str]:
     if isinstance(obj, DataArray):
         obj = obj._to_temp_dataset()
 
+    variables = obj._variables
     results = set()
-    for var in obj.variables:
-        da = obj[var]
-        attrs_or_encoding = ChainMap(da.attrs, da.encoding)
+    for var in variables.values():
+        attrs_or_encoding = ChainMap(var.attrs, var.encoding)
         if "grid_mapping" in attrs_or_encoding:
             grid_mapping_var_name = attrs_or_encoding["grid_mapping"]
-            if grid_mapping_var_name not in obj.variables:
+            if grid_mapping_var_name not in variables:
                 raise ValueError(
                     f"{var} defines non-existing grid_mapping variable {grid_mapping_var_name}."
                 )
-            if key == obj[grid_mapping_var_name].attrs["grid_mapping_name"]:
+            if key == variables[grid_mapping_var_name].attrs["grid_mapping_name"]:
                 results.update([grid_mapping_var_name])
     return list(results)
 
@@ -474,7 +481,7 @@ def _get_indexes(obj: DataArray | Dataset, key: Hashable) -> list[Hashable]:
     One or more of ('X', 'Y', 'Z', 'T', 'longitude', 'latitude', 'vertical', 'time',
     'area', 'volume'), or arbitrary measures, or standard names present in .indexes
     """
-    return [k for k in _get_all(obj, key) if k in obj.indexes]
+    return [k for k in _get_all(obj, key) if k in obj._indexes]
 
 
 def _get_coords(obj: DataArray | Dataset, key: Hashable) -> list[Hashable]:
@@ -2251,7 +2258,7 @@ def get_bounds(self, key: Hashable) -> DataArray | Dataset:
         DataArray
         """
 
-        results = self.bounds.get(key, [])
+        results = self[[key]].cf.bounds.get(key, [])
         if not results:
             raise KeyError(f"No results found for {key!r}.")
 
@@ -2270,12 +2277,18 @@ def get_bounds_dim_name(self, key: Hashable) -> Hashable:
         -------
         str
         """
-        crd = self[key]
-        bounds = self.get_bounds(key)
+        (crd_name,) = apply_mapper(_get_all, self._obj, key, error=False, default=[key])
+        variables = self._obj._variables
+        crd = variables[crd_name]
+        crd_attrs = crd._attrs
+        if crd_attrs is None or "bounds" not in crd_attrs:
+            raise KeyError(f"No bounds variable found for {key!r}")
+
+        bounds = variables[crd_attrs["bounds"].strip()]
         bounds_dims = set(bounds.dims) - set(crd.dims)
         assert len(bounds_dims) == 1
         bounds_dim = bounds_dims.pop()
-        assert self._obj.sizes[bounds_dim] in [2, 4]
+        assert bounds.sizes[bounds_dim] in [2, 4]
         return bounds_dim
 
     def add_bounds(