feat: created queries for talk to ipcc

climateqa/engine/talk_to_data/ipcc/queries.py

@@ -0,0 +1,116 @@
+from typing import TypedDict, Optional
+
+from climateqa.engine.talk_to_data.ipcc.config import IPCC_DATASET_URL, MACRO_COUNTRIES
+
+class IndicatorPerYearAtLocationQueryParams(TypedDict, total=False):
+    """
+    Parameters for querying the evolution of an indicator per year at a specific location.
+
+    Attributes:
+        indicator_column (str): Name of the climate indicator column.
+        latitude (str): Latitude of the location.
+        longitude (str): Longitude of the location.
+        country_code (str): Country code.
+        admin1 (str): Administrative region (optional).
+    """
+    indicator_column: str
+    latitude: str
+    longitude: str
+    country_code: str
+    admin1: Optional[str]
+
+def indicator_per_year_at_location_query(
+    table: str, params: IndicatorPerYearAtLocationQueryParams
+) -> str:
+    """
+    Builds an SQL query to get the evolution of an indicator per year at a specific location.
+
+    Args:
+        table (str): SQL table of the indicator.
+        params (IndicatorPerYearAtLocationQueryParams): Dictionary with the required params for the query.
+
+    Returns:
+        str: The SQL query string, or an empty string if required parameters are missing.
+    """
+    indicator_column = params.get("indicator_column")
+    latitude = params.get("latitude")
+    longitude = params.get("longitude")
+    country_code = params.get("country_code")
+    admin1 = params.get("admin1")
+
+    if not all([indicator_column, latitude, longitude, country_code]):
+        return ""
+
+    if country_code in MACRO_COUNTRIES:
+        table_path = f"'{IPCC_DATASET_URL}/{table.lower()}/{country_code}_macro.parquet'"
+        sql_query = f"""
+        SELECT year, scenario, {indicator_column}
+        FROM {table_path}
+        WHERE admin1 = '{admin1}' AND year >= 1950
+        ORDER BY year, scenario
+        """
+    else:
+        table_path = f"'{IPCC_DATASET_URL}/{table.lower()}/{country_code}.parquet'"
+        sql_query = f"""
+        SELECT year, scenario, AVG({indicator_column}) AS {indicator_column}
+        FROM {table_path}
+        WHERE latitude = {latitude} AND longitude = {longitude} AND year >= 1950
+        GROUP BY scenario, year
+        ORDER BY year, scenario
+        """
+    return sql_query.strip()
+
+class IndicatorForGivenYearQueryParams(TypedDict, total=False):
+    """
+    Parameters for querying an indicator's values across locations for a specific year.
+
+    Attributes:
+        indicator_column (str): The column name for the climate indicator.
+        year (str): The year to query.
+        country_code (str): The country code.
+    """
+    indicator_column: str
+    year: str
+    country_code: str
+
+def indicator_for_given_year_query(
+    table: str, params: IndicatorForGivenYearQueryParams
+) -> str:
+    """
+    Builds an SQL query to get the values of an indicator with their latitudes, longitudes,
+    and scenarios for a given year.
+
+    Args:
+        table (str): SQL table of the indicator.
+        params (IndicatorForGivenYearQueryParams): Dictionary with the required params for the query.
+
+    Returns:
+        str: The SQL query string, or an empty string if required parameters are missing.
+    """
+    indicator_column = params.get("indicator_column")
+    year = params.get("year") or 2050
+    country_code = params.get("country_code")
+
+    if not all([indicator_column, year, country_code]):
+        return ""
+
+    if country_code in MACRO_COUNTRIES:
+        table_path = f"'{IPCC_DATASET_URL}/{table.lower()}/{country_code}_macro.parquet'"
+        sql_query = f"""
+        SELECT {indicator_column}, c.latitude, c.longitude, c.admin1, scenario
+        FROM {table_path} AS t
+        RIGHT JOIN '{IPCC_DATASET_URL}/coordinates.parquet' AS c 
+        ON c.admin1 = t.admin1 AND c.country_code = t.country_code
+        WHERE year = {year}
+        ORDER BY latitude, longitude, scenario
+        """
+    else:
+        table_path = f"'{IPCC_DATASET_URL}/{table.lower()}/{country_code}.parquet'"
+        sql_query = f"""
+        SELECT AVG({indicator_column}) AS {indicator_column}, latitude, longitude, scenario
+        FROM {table_path}
+        WHERE year = {year}
+        GROUP BY latitude, longitude, scenario
+        ORDER BY latitude, longitude, scenario
+        """
+    return sql_query.strip()