feature: support property filters in observations_dataframe (#225)

Feature: specify property filters so that only the relevant observations are fetched. This could substantially reduce the
amount of data transferred if users specify a particular `unit`,`measurementMethod`, `observationPeriod`, etc.

Author: jm-rivera

datacommons_client/client.py

@@ -6,6 +6,7 @@
 from datacommons_client.endpoints.payloads import ObservationDate
 from datacommons_client.endpoints.resolve import ResolveEndpoint
 from datacommons_client.utils.decorators import requires_pandas
+from datacommons_client.utils.error_handling import NoDataForPropertyError
 
 try:
   import pandas as pd
@@ -58,6 +59,58 @@ def __init__(
     self.observation = ObservationEndpoint(api=self.api)
     self.resolve = ResolveEndpoint(api=self.api)
 
+  def _find_filter_facet_ids(
+      self,
+      fetch_by: Literal["entity", "entity_type"],
+      date: ObservationDate | str,
+      variable_dcids: str | list[str],
+      entity_dcids: Literal["all"] | list[str] = "all",
+      entity_type: Optional[str] = None,
+      parent_entity: Optional[str] = None,
+      property_filters: Optional[dict[str, str | list[str]]] = None,
+  ) -> list[str] | None:
+    """Finds matching facet IDs for property filters.
+
+        Args:
+            fetch_by (Literal["entity", "entity_type"]): Determines whether to fetch by entity or entity type.
+            variable_dcids (str | list[str]): The variable DCIDs for which to retrieve facet IDs.
+            entity_dcids (Literal["all"] | list[str], optional): The entity DCIDs, or "all" if filtering by entity type.
+            entity_type (Optional[str]): The entity type, required if fetching by entity type.
+            parent_entity (Optional[str]): The parent entity, used when fetching by entity type.
+            property_filters (Optional[dict[str, str | list[str]]): A dictionary of properties to match facets against.
+
+        Returns:
+            list[str] | None: A list of matching facet IDs, or None if no filters are applied.
+        """
+
+    if not property_filters:
+      return None
+
+    if fetch_by == "entity":
+      observations = self.observation.fetch_observations_by_entity(
+          date=date,
+          entity_dcids=entity_dcids,
+          variable_dcids=variable_dcids,
+          select=["variable", "entity", "facet"],
+      )
+    else:
+      observations = self.observation.fetch_observations_by_entity_type(
+          date=date,
+          entity_type=entity_type,
+          parent_entity=parent_entity,
+          variable_dcids=variable_dcids,
+          select=["variable", "entity", "facet"],
+      )
+
+    facet_sets = [
+        observations.find_matching_facet_id(property_name=p, value=v)
+        for p, v in property_filters.items()
+    ]
+
+    facet_ids = list({facet for facets in facet_sets for facet in facets})
+
+    return facet_ids
+
   @requires_pandas
   def observations_dataframe(
       self,
@@ -66,6 +119,7 @@ def observations_dataframe(
       entity_dcids: Literal["all"] | list[str] = "all",
       entity_type: Optional[str] = None,
       parent_entity: Optional[str] = None,
+      property_filters: Optional[dict[str, str | list[str]]] = None,
   ):
     """
         Fetches statistical observations and returns them as a Pandas DataFrame.
@@ -74,15 +128,17 @@ def observations_dataframe(
         at a particular date (e.g., "population of USA in 2020", "GDP of California in 2010").
 
         Args:
-            variable_dcids (str | list[str]): One or more variable DCIDs for the observation.
-            date (ObservationDate | str): The date for which observations are requested. It can be
+        variable_dcids (str | list[str]): One or more variable DCIDs for the observation.
+        date (ObservationDate | str): The date for which observations are requested. It can be
             a specific date, "all" to retrieve all observations, or "latest" to get the most recent observations.
-            entity_dcids (Literal["all"] | list[str], optional): The entity DCIDs to retrieve data for.
-                Defaults to "all". DCIDs must include their type (e.g "country/GTM" for Guatemala).
-            entity_type (Optional[str], optional): The type of entities to filter by when `entity_dcids="all"`.
-                Required if `entity_dcids="all"`. Defaults to None.
-            parent_entity (Optional[str], optional): The parent entity under which the target entities fall.
-                Used only when `entity_dcids="all"`. Defaults to None.
+        entity_dcids (Literal["all"] | list[str], optional): The entity DCIDs to retrieve data for.
+            Defaults to "all". DCIDs must include their type (e.g., "country/GTM" for Guatemala).
+        entity_type (Optional[str]): The type of entities to filter by when `entity_dcids="all"`.
+            Required if `entity_dcids="all"`. Defaults to None.
+        parent_entity (Optional[str]): The parent entity under which the target entities fall.
+            Used only when `entity_dcids="all"`. Defaults to None.
+        property_filters (Optional[dict[str, str | list[str]]): An optional dictionary used to filter
+            the data by using observation properties like `measurementMethod`, `unit`, or `observationPeriod`.
 
         Returns:
             pd.DataFrame: A DataFrame containing the requested observations.
@@ -97,14 +153,34 @@ def observations_dataframe(
           "Specify 'entity_type' and 'parent_entity' only when 'entity_dcids' is 'all'."
       )
 
+    # If property filters are provided, fetch the required facet IDs. Otherwise, set to None.
+    facets = self._find_filter_facet_ids(
+        fetch_by="entity" if entity_dcids != "all" else "entity_type",
+        date=date,
+        variable_dcids=variable_dcids,
+        entity_dcids=entity_dcids,
+        entity_type=entity_type,
+        parent_entity=parent_entity,
+        property_filters=property_filters,
+    )
+
+    if not facets and property_filters:
+      raise NoDataForPropertyError
+
     if entity_dcids == "all":
       observations = self.observation.fetch_observations_by_entity_type(
           date=date,
           parent_entity=parent_entity,
           entity_type=entity_type,
-          variable_dcids=variable_dcids)
+          variable_dcids=variable_dcids,
+          filter_facet_ids=facets,
+      )
     else:
       observations = self.observation.fetch_observations_by_entity(
-          date=date, entity_dcids=entity_dcids, variable_dcids=variable_dcids)
+          date=date,
+          entity_dcids=entity_dcids,
+          variable_dcids=variable_dcids,
+          filter_facet_ids=facets,
+      )
 
     return pd.DataFrame(observations.get_observations_as_records())

datacommons_client/endpoints/observation.py

@@ -68,6 +68,7 @@ def fetch_latest_observations(
       entity_dcids: Optional[str | list[str]] = None,
       entity_expression: Optional[str] = None,
       *,
+      select: Optional[list[ObservationSelect | str]] = None,
       filter_facet_domains: Optional[str | list[str]] = None,
       filter_facet_ids: Optional[str | list[str]] = None,
   ) -> ObservationResponse:
@@ -78,6 +79,8 @@ def fetch_latest_observations(
             variable_dcids (str | list[str]): One or more variable IDs for the data.
             entity_dcids (Optional[str | list[str]]): One or more entity IDs to filter the data.
             entity_expression (Optional[str]): A string expression to filter entities.
+            select (Optional[list[ObservationSelect | str]]): Fields to include in the response.
+                If not provided, defaults to ["date", "variable", "entity", "value"].
             filter_facet_domains: Optional[str | list[str]: One or more domain names to filter the data.
             filter_facet_ids: Optional[str | list[str]: One or more facet IDs to filter the data.
 
@@ -91,13 +94,15 @@ def fetch_latest_observations(
         entity_expression=entity_expression,
         filter_facet_domains=filter_facet_domains,
         filter_facet_ids=filter_facet_ids,
+        select=[s for s in ObservationSelect] if not select else select,
     )
 
   def fetch_latest_observations_by_entity(
       self,
       variable_dcids: str | list[str],
       entity_dcids: str | list[str],
       *,
+      select: Optional[list[ObservationSelect | str]] = None,
       filter_facet_domains: Optional[str | list[str]] = None,
       filter_facet_ids: Optional[str | list[str]] = None,
   ) -> ObservationResponse:
@@ -106,6 +111,8 @@ def fetch_latest_observations_by_entity(
         Args:
             variable_dcids (str | list[str]): One or more variable IDs for the data.
             entity_dcids (str | list[str]): One or more entity IDs to filter the data.
+            select (Optional[list[ObservationSelect | str]]): Fields to include in the response.
+                If not provided, defaults to ["date", "variable", "entity", "value"].
             filter_facet_domains: Optional[str | list[str]: One or more domain names to filter the data.
             filter_facet_ids: Optional[str | list[str]: One or more facet IDs to filter the data.
 
@@ -116,8 +123,10 @@ def fetch_latest_observations_by_entity(
     return self.fetch_latest_observations(
         variable_dcids=variable_dcids,
         entity_dcids=entity_dcids,
+        select=[s for s in ObservationSelect] if not select else select,
         filter_facet_domains=filter_facet_domains,
-        filter_facet_ids=filter_facet_ids)
+        filter_facet_ids=filter_facet_ids,
+    )
 
   def fetch_observations_by_entity_type(
       self,
@@ -126,6 +135,7 @@ def fetch_observations_by_entity_type(
       entity_type: str,
       variable_dcids: str | list[str],
       *,
+      select: Optional[list[ObservationSelect | str]] = None,
       filter_facet_domains: Optional[str | list[str]] = None,
       filter_facet_ids: Optional[str | list[str]] = None,
   ) -> ObservationResponse:
@@ -142,6 +152,8 @@ def fetch_observations_by_entity_type(
                 For example, "Country" or "Region".
             variable_dcids (str | list[str]): The variable(s) to fetch observations for.
                 This can be a single variable ID or a list of IDs.
+            select (Optional[list[ObservationSelect | str]]): Fields to include in the response.
+                If not provided, defaults to ["date", "variable", "entity", "value"].
             filter_facet_domains: Optional[str | list[str]: One or more domain names to filter the data.
             filter_facet_ids: Optional[str | list[str]: One or more facet IDs to filter the data.
 
@@ -165,7 +177,7 @@ def fetch_observations_by_entity_type(
     return self.fetch(
         variable_dcids=variable_dcids,
         date=date,
-        select=[s for s in ObservationSelect],
+        select=[s for s in ObservationSelect] if not select else select,
         entity_expression=
         f"{parent_entity}<-containedInPlace+{{typeOf:{entity_type}}}",
         filter_facet_domains=filter_facet_domains,
@@ -178,6 +190,7 @@ def fetch_observations_by_entity(
       entity_dcids: str | list[str],
       variable_dcids: str | list[str],
       *,
+      select: Optional[list[ObservationSelect | str]] = None,
       filter_facet_domains: Optional[str | list[str]] = None,
       filter_facet_ids: Optional[str | list[str]] = None,
   ) -> ObservationResponse:
@@ -191,6 +204,8 @@ def fetch_observations_by_entity(
             entity_dcids (str | list[str]): One or more entity IDs to filter the data.
             variable_dcids (str | list[str]): The variable(s) to fetch observations for.
                 This can be a single variable ID or a list of IDs.
+            select (Optional[list[ObservationSelect | str]]): Fields to include in the response.
+                If not provided, defaults to ["date", "variable", "entity", "value"].
             filter_facet_domains: Optional[str | list[str]: One or more domain names to filter the data.
             filter_facet_ids: Optional[str | list[str]: One or more facet IDs to filter the data.
 
@@ -213,7 +228,7 @@ def fetch_observations_by_entity(
     return self.fetch(
         variable_dcids=variable_dcids,
         date=date,
-        select=[s for s in ObservationSelect],
+        select=[s for s in ObservationSelect] if not select else select,
         entity_dcids=entity_dcids,
         filter_facet_domains=filter_facet_domains,
         filter_facet_ids=filter_facet_ids,

datacommons_client/endpoints/payloads.py

@@ -78,6 +78,7 @@ class ObservationSelect(str, Enum):
   VARIABLE = "variable"
   ENTITY = "entity"
   VALUE = "value"
+  FACET = "facet"
 
   @classmethod
   def _missing_(cls, value):
@@ -138,7 +139,6 @@ def __post_init__(self):
   def normalize(self):
     """
         Normalizes the payload for consistent internal representation.
-
         - Converts `variable_dcids`, `entity_dcids`, `filter_facet_domains` and `filter_facet_ids`
          to lists if they are passed as strings.
         - Normalizes the `date` field to ensure it is in the correct format.

datacommons_client/endpoints/response.py

@@ -151,6 +151,79 @@ def get_observations_as_records(self) -> List[Dict[str, Any]]:
     return observations_as_records(data=self.get_data_by_entity(),
                                    facets=self.facets)
 
+  def get_facets_metadata(self) -> Dict[str, Any]:
+    """Extract metadata about StatVars from the response. This data is
+        structured as a dictionary of StatVars, each containing a dictionary of
+        facets with their corresponding metadata.
+
+        Returns:
+            Dict[str, Any]: A dictionary of StatVars with their associated metadata,
+             including earliest and latest observation dates, observation counts,
+             measurementMethod, observationPeriod, and unit, etc.
+        """
+    # Dictionary to store metadata
+    metadata = {}
+
+    # Extract information from byVariable
+    data_by_entity = self.get_data_by_entity()
+
+    # Extract facet information
+    facets_info = self.to_dict().get("facets", {})
+
+    for dcid, variables in data_by_entity.items():
+      metadata[dcid] = {}
+
+      for entity_id, entity in variables.items():
+        for facet in entity.get("orderedFacets", []):
+          facet_metadata = metadata[dcid].setdefault(
+              facet.facetId,
+              {
+                  "earliestDate": {},
+                  "latestDate": {},
+                  "obsCount": {},
+              },
+          )
+
+          facet_metadata["earliestDate"][entity_id] = facet.earliestDate
+          facet_metadata["latestDate"][entity_id] = facet.latestDate
+          facet_metadata["obsCount"][entity_id] = facet.obsCount
+
+          # Merge additional facet details
+          facet_metadata.update(facets_info.get(facet.facetId, {}))
+
+    return metadata
+
+  def find_matching_facet_id(self, property_name: str,
+                             value: str | list[str]) -> list[str]:
+    """Finds facet IDs that match a given property and value.
+
+        Args:
+            property_name (str): The property to match.
+            value (str | list): The value to match. Can be a string, number, or a list of values.
+        Returns:
+            list[str]: A list of facet IDs that match the property and value.
+        """
+    # Initialize an empty list to store matching facet IDs
+    matching_facet_ids = []
+
+    # Iterate over the facets metadata to find matching facet IDs
+    for facet_data in self.get_facets_metadata().values():
+
+      # Iterate over each facet and its associated metadata
+      for facet_id, metadata in facet_data.items():
+
+        # Get the value of the specified property from the data
+        prop_value = metadata.get(property_name)
+
+        # Check if the property value matches the specified value
+        if isinstance(value, list):
+          if prop_value in value:
+            matching_facet_ids.append(facet_id)
+        elif prop_value == value:
+          matching_facet_ids.append(facet_id)
+
+    return matching_facet_ids
+
 
 @dataclass
 class ResolveResponse(SerializableMixin):

datacommons_client/tests/endpoints/test_error_handling.py

@@ -7,6 +7,7 @@
 from datacommons_client.utils.error_handling import DCConnectionError
 from datacommons_client.utils.error_handling import DCStatusError
 from datacommons_client.utils.error_handling import InvalidDCInstanceError
+from datacommons_client.utils.error_handling import NoDataForPropertyError
 
 
 def test_data_commons_error_default_message():
@@ -59,6 +60,9 @@ def test_subclass_default_messages():
   instance_error = InvalidDCInstanceError()
   assert InvalidDCInstanceError.default_message in str(instance_error)
 
+  filter_error = NoDataForPropertyError()
+  assert NoDataForPropertyError.default_message in str(filter_error)
+
 
 def test_subclass_custom_message():
   """Tests that subclasses use custom messages when provided."""