LAMA to Dask: fixes to Query & Data.digitize

cf/data/data.py

@@ -1787,6 +1787,18 @@ def digitize(
         d._set_dask(dx)
         d.override_units(_units_None, inplace=True)
 
+        # More elegant to handle 'delete_bins' in cf- rather than Dask- space
+        # i.e. using cf.where with d in-place rather than da.where with dx
+        # just after the digitize operation above (cf.where already applies
+        # equivalent logic element-wise).
+        if delete_bins:
+            for n, db in enumerate(delete_bins):
+                db -= n
+                d.where(d == db, np.ma.masked, None, inplace=True)
+                # x = d - 1 rather than = d here since there is one fewer bin
+                # therefore we need to adjust to the new corresponding indices
+                d.where(d > db, d - 1, None, inplace=True)
+
         if return_bins:
             if two_d_bins is None:
                 two_d_bins = np.empty((bins.size - 1, 2), dtype=bins.dtype)

cf/query.py

@@ -857,32 +857,47 @@ def set_condition_units(self, units):
         <CF Query: (ge 3000 m)>
 
         """
+        def get_and_set_value_units(v, u):
+            """Helper method to simplify logic to set specified units."""
+            v_units = getattr(v, "Units", None)
+            if v_units is None:  # Value 'v' has no units
+                v = Data(v, units=u)
+            else:  # Value 'v' already has units
+                try:
+                    v = v.copy()
+                    v.Units = u
+                except ValueError:
+                    raise ValueError(
+                        f"Units {u!r} are not equivalent to "
+                        f"query condition units {v_units!r}"
+                    )
+
+            return v
+
         units = Units(units)
 
         compound = self._compound
         if compound:
             for r in compound:
                 r.set_condition_units(units)
-
             return
 
         value = self._value
         if value is None:
             return
 
         value_units = getattr(value, "Units", None)
-        if value_units is None:
-            # Value has no units
-            value = Data(value, units=units)
+        if self.operator in ("wi", "wo", "set"):
+            # Value is a sequence of things that may or may not
+            # already have units
+            new_values = []
+            for v in value:
+                v = get_and_set_value_units(v, units)
+                new_values.append(v)
+
+            value = new_values
         else:
-            # Value already has units
-            try:
-                value.Units = units
-            except ValueError:
-                raise ValueError(
-                    f"Units {units!r} are not equivalent to "
-                    f"query condition units {value_units!r}"
-                )
+            value = get_and_set_value_units(value, units)
 
         self._value = value
 

cf/test/test_Data.py

@@ -942,18 +942,14 @@ def test_Data_digitize(self):
                         (np.ma.getmask(e.array) == np.ma.getmask(b)).all()
                     )
 
-                    # TODODASK: Reinstate the following test when
-                    #           __sub__, minimum, and maximum have
-                    #           been daskified
-
-        #                    e.where(
-        #                        cf.set([e.minimum(), e.maximum()]),
-        #                        cf.masked,
-        #                        e - 1,
-        #                        inplace=True,
-        #                    )
-        #                    f = d.digitize(bins, upper=upper)
-        #                    self.assertTrue(e.equals(f, verbose=2))
+                    e.where(
+                        cf.set([e.minimum(), e.maximum()]),
+                        cf.masked,
+                        e - 1,
+                        inplace=True,
+                    )
+                    f = d.digitize(bins, upper=upper)
+                    self.assertTrue(e.equals(f, verbose=2))
 
         # Check returned bins
         bins = [2, 6, 10, 50, 100]