Improve search with derived variables

intake_esm/core.py

@@ -329,13 +329,25 @@ def search(
         # step 2: Search for entries required to derive variables in the derived catalogs
         # This requires a bit of a hack i.e. the user has to specify the variable in the query
         derivedcat_results = []
-        variables = query.get(self.esmcat.aggregation_control.variable_column_name, [])
+        variables = query.pop(self.esmcat.aggregation_control.variable_column_name, None)
         if variables:
+            if isinstance(variables, str):
+                variables = [variables]
+            dependents = []
             for key, value in self.derivedcat.items():
                 if key in variables:
-                    derivedcat_results.append(
-                        self.esmcat.search(require_all_on=require_all_on, query=value.query)
+                    res = self.esmcat.search(
+                        require_all_on=require_all_on, query={**value.query, **query}
                     )
+                    if not res.empty:
+                        derivedcat_results.append(res)
+                        dependents.extend(
+                            value.dependent_variables(
+                                self.esmcat.aggregation_control.variable_column_name
+                            )
+                        )
+            # Update variable list with dependent variables.
+            variables = list(set(variables).union(dependents))
 
         if derivedcat_results:
             # Merge results from the main and the derived catalogs
@@ -347,16 +359,18 @@ def search(
 
         cat = self.__class__({'esmcat': self.esmcat.dict(), 'df': esmcat_results})
         if self.esmcat.has_multiple_variable_assets:
-            requested_variables = query.get(
-                self.esmcat.aggregation_control.variable_column_name, []
-            )
+            requested_variables = variables or []
         else:
             requested_variables = []
         cat._requested_variables = requested_variables
 
-        # step 3: Subset the derived catalog
-        derivat_cat_subset = self.derivedcat.search(variable=variables)
-        cat.derivedcat = derivat_cat_subset
+        # step 3: Subset the derived catalog,
+        # but only if variables were looked up, otherwise transfer the whole catalog.
+        if variables is not None:
+            derivat_cat_subset = self.derivedcat.search(variable=variables)
+            cat.derivedcat = derivat_cat_subset
+        else:
+            cat.derivedcat = self.derivedcat
         return cat
 
     @pydantic.validate_arguments

tests/test_core.py

@@ -22,6 +22,14 @@ def funcs(ds):
     return ds + 1
 
 
+registry_multivar = intake_esm.DerivedVariableRegistry()
+
+
+@registry_multivar.register(variable='FOO', query={'variable': ['TEMP']})
+def func_multivar(ds):
+    return ds + 1
+
+
 from .utils import (
     catalog_dict_records,
     cdf_col_sample_cesmle,
@@ -125,14 +133,21 @@ def test_catalog_search(path, query, expected_size):
 
 
 def test_catalog_with_registry_search():
-    cat = intake.open_esm_datastore(catalog_dict_records, registry=registry)
+    cat = intake.open_esm_datastore(zarr_col_aws_cesm, registry=registry)
     new_cat = cat.search(variable='FOO')
-    assert len(cat) == 1
-    assert len(new_cat) == 1
+    assert len(cat) == 56
+    assert len(new_cat) == 11
 
     assert len(cat.derivedcat) == 2
     assert len(new_cat.derivedcat) == 1
 
+    new_cat = cat.search(variable='FOO', frequency='daily')
+    assert len(new_cat) == 4
+    assert len(new_cat.derivedcat) == 1
+
+    new_cat = cat.search(frequency='daily')
+    assert len(new_cat.derivedcat) == 2
+
 
 @pytest.mark.parametrize('key', ['ocn.20C.pop.h', 'ocn.CTRL.pop.h', 'ocn.RCP85.pop.h'])
 def test_catalog_getitem(key):
@@ -221,6 +236,18 @@ def test_multi_variable_catalog(query):
         assert set(ds.data_vars) == set(cat_sub._requested_variables)
 
 
+def test_multi_variable_catalog_derived_cat():
+    import ast
+
+    cat = intake.open_esm_datastore(
+        multi_variable_col,
+        read_csv_kwargs={'converters': {'variable': ast.literal_eval}},
+        registry=registry_multivar,
+    )
+    cat_sub = cat.search(variable=['FOO'])
+    assert set(cat_sub._requested_variables) == {'TEMP', 'FOO'}
+
+
 @pytest.mark.parametrize(
     'path, query, xarray_open_kwargs',
     [