Server-side compute: stubs and flags

Author: dshemetov

requirements.txt

@@ -1,5 +1,6 @@
 Flask==1.1.2
 SQLAlchemy==1.3.22
+more_itertools==8.4.0
 mysqlclient==2.0.2
 python-dotenv==0.15.0
 orjson==3.4.7
@@ -8,3 +9,4 @@ scipy==1.6.2
 tenacity==7.0.0
 newrelic
 epiweeks==2.1.2
+delphi_utils

src/server/_params.py

@@ -1,11 +1,11 @@
+from dataclasses import dataclass
+from itertools import groupby, chain
 from math import inf
 import re
-from dataclasses import dataclass
 from typing import List, Optional, Sequence, Tuple, Union
 
 from flask import request
 
-
 from ._exceptions import ValidationFailedException
 from .utils import days_in_range, weeks_in_range, guess_time_value_is_day
 
@@ -93,8 +93,17 @@ def count(self) -> float:
         return len(self.signal)
 
 
+def _combine_source_signal_pairs(source_signal_pairs: List[SourceSignalPair]) -> List[SourceSignalPair]:
+    """Combine SourceSignalPairs with the same source into a single SourceSignalPair object.
+
+    Example:
+    [SourceSignalPair("src", ["sig1", "sig2"]), SourceSignalPair("src", ["sig2", "sig3"])] will be merged
+    into [SourceSignalPair("src", ["sig1", "sig2", "sig3])].
+    """
+    return [SourceSignalPair("src", ["sig1", "sig2", "sig3"])]
+
 def parse_source_signal_arg(key: str = "signal") -> List[SourceSignalPair]:
-    return [SourceSignalPair(source, signals) for [source, signals] in _parse_common_multi_arg(key)]
+    return _combine_source_signal_pairs([SourceSignalPair(source, signals) for [source, signals] in _parse_common_multi_arg(key)])
 
 
 def parse_single_source_signal_arg(key: str) -> SourceSignalPair:

src/server/endpoints/covidcast.py

@@ -1,10 +1,11 @@
 from typing import List, Optional, Union, Tuple, Dict, Any
 from itertools import groupby
 from datetime import date, timedelta
+from bisect import bisect_right
 from epiweeks import Week
 from flask import Blueprint, request
 from flask.json import loads, jsonify
-from bisect import bisect_right
+from more_itertools import peekable
 from sqlalchemy import text
 from pandas import read_csv, to_datetime
 
@@ -36,11 +37,13 @@
 from .._pandas import as_pandas, print_pandas
 from .covidcast_utils import compute_trend, compute_trends, compute_correlations, compute_trend_value, CovidcastMetaEntry
 from ..utils import shift_time_value, date_to_time_value, time_value_to_iso, time_value_to_date, shift_week_value, week_value_to_week, guess_time_value_is_day, week_to_time_value
-from .covidcast_utils.model import TimeType, count_signal_time_types, data_sources, create_source_signal_alias_mapper
+from .covidcast_utils.model import TimeType, count_signal_time_types, data_sources, create_source_signal_alias_mapper, get_basename_signals, get_pad_length, pad_time_pairs, pad_time_window
+from .covidcast_utils.smooth_diff import PadFillValue, SmootherKernelValue
 
 # first argument is the endpoint name
 bp = Blueprint("covidcast", __name__)
 alias = None
+JIT_COMPUTE = True
 
 
 def parse_source_signal_pairs() -> List[SourceSignalPair]:
@@ -123,36 +126,53 @@ def guess_index_to_use(time: List[TimePair], geo: List[GeoPair], issues: Optiona
     return None
 
 
+# TODO: Write an actual smoother arg parser.
+def parse_transform_args():
+    return {"smoother_kernel": SmootherKernelValue.average, "smoother_window_length": 7, "pad_fill_value": None, "nans_fill_value": None}
+
+
+def parse_jit_bypass():
+    jit_bypass = request.values.get("jit_bypass")
+    if jit_bypass is not None:
+        return bool(jit_bypass)
+    else:
+        return False
+
+
 @bp.route("/", methods=("GET", "POST"))
 def handle():
     source_signal_pairs = parse_source_signal_pairs()
     source_signal_pairs, alias_mapper = create_source_signal_alias_mapper(source_signal_pairs)
     time_pairs = parse_time_pairs()
     geo_pairs = parse_geo_pairs()
+    transform_args = parse_transform_args()
+    jit_bypass = parse_jit_bypass()
 
     as_of = extract_date("as_of")
     issues = extract_dates("issues")
     lag = extract_integer("lag")
+    is_time_type_week = any(time_pair.time_type == "week" for time_pair in time_pairs)
+    is_time_value_true = any(isinstance(time_pair.time_values, bool) for time_pair in time_pairs)
+    use_server_side_compute = not any((issues, lag, is_time_type_week, is_time_value_true)) and JIT_COMPUTE and not jit_bypass
+    if use_server_side_compute:
+        pad_length = get_pad_length(source_signal_pairs, transform_args.get("smoother_window_length"))
+        source_signal_pairs, row_transform_generator = get_basename_signals(source_signal_pairs)
+        time_pairs = pad_time_pairs(time_pairs, pad_length)
 
     # build query
     q = QueryBuilder("covidcast", "t")
 
-    fields_string = ["geo_value", "signal"]
+    fields_string = ["geo_type", "geo_value", "source", "signal", "time_type"]
     fields_int = ["time_value", "direction", "issue", "lag", "missing_value", "missing_stderr", "missing_sample_size"]
     fields_float = ["value", "stderr", "sample_size"]
     is_compatibility = is_compatibility_mode()
-    if is_compatibility:
-        q.set_order("signal", "time_value", "geo_value", "issue")
-    else:
-        # transfer also the new detail columns
-        fields_string.extend(["source", "geo_type", "time_type"])
-        q.set_order("source", "signal", "time_type", "time_value", "geo_type", "geo_value", "issue")
+
+    q.set_order("geo_type", "geo_value", "source", "signal", "time_type", "time_value", "issue")
     q.set_fields(fields_string, fields_int, fields_float)
 
     # basic query info
     # data type of each field
     # build the source, signal, time, and location (type and id) filters
-
     q.where_source_signal_pairs("source", "signal", source_signal_pairs)
     q.where_geo_pairs("geo_type", "geo_value", geo_pairs)
     q.where_time_pairs("time_type", "time_value", time_pairs)
@@ -161,14 +181,40 @@ def handle():
 
     _handle_lag_issues_as_of(q, issues, lag, as_of)
 
-    def transform_row(row, proxy):
+    p = create_printer()
+
+    def alias_row(row):
+        if is_compatibility:
+            # old api returned fewer fields
+            remove_fields = ["geo_type", "source", "time_type"]
+            for field in remove_fields:
+                if field in row:
+                    del row[field]
         if is_compatibility or not alias_mapper or "source" not in row:
             return row
-        row["source"] = alias_mapper(row["source"], proxy["signal"])
+        row["source"] = alias_mapper(row["source"], row["signal"])
         return row
 
-    # send query
-    return execute_query(str(q), q.params, fields_string, fields_int, fields_float, transform=transform_row)
+    if use_server_side_compute:
+        def gen_transform(rows):
+            parsed_rows = (parse_row(row, fields_string, fields_int, fields_float) for row in rows)
+            transformed_rows = row_transform_generator(parsed_rows=parsed_rows, time_pairs=time_pairs, transform_args=transform_args)
+            for row in transformed_rows:
+                yield alias_row(row)
+    else:
+        def gen_transform(rows):
+            parsed_rows = (parse_row(row, fields_string, fields_int, fields_float) for row in rows)
+            for row in parsed_rows:
+                yield alias_row(row)
+
+    # execute first query
+    try:
+        r = run_query(p, (str(q), q.params))
+    except Exception as e:
+        raise DatabaseErrorException(str(e))
+
+    # now use a generator for sending the rows and execute all the other queries
+    return p(filter_fields(gen_transform(r)))
 
 
 def _verify_argument_time_type_matches(is_day_argument: bool, count_daily_signal: int, count_weekly_signal: int) -> None:

src/server/endpoints/covidcast_utils/model.py

@@ -1,12 +1,21 @@
 from dataclasses import asdict, dataclass, field
-from typing import Callable, Optional, Dict, List, Set, Tuple
 from enum import Enum
+from functools import partial
+from itertools import groupby, repeat, tee
+from typing import Callable, Generator, Optional, Dict, List, Set, Tuple, Iterable, Union
 from pathlib import Path
 import re
+from more_itertools import interleave_longest, peekable
 import pandas as pd
 import numpy as np
 
-from ..._params import SourceSignalPair
+from delphi_utils.nancodes import Nans
+from ..._params import SourceSignalPair, TimePair
+from .smooth_diff import generate_smooth_rows, generate_row_diffs
+from ...utils import time_value_range, shift_time_value
+
+
+IDENTITY: Callable = lambda rows, **kwargs: rows
 
 
 class HighValuesAre(str, Enum):
@@ -21,6 +30,7 @@ class SignalFormat(str, Enum):
     fraction = "fraction"
     raw_count = "raw_count"
     raw = "raw"
+    count = "count"
 
 
 class SignalCategory(str, Enum):
@@ -201,7 +211,7 @@ def _load_data_sources():
 
 
 data_sources, data_sources_df = _load_data_sources()
-data_source_by_id = {d.source: d for d in data_sources}
+data_sources_by_id = {d.source: d for d in data_sources}
 
 
 def _load_data_signals(sources: List[DataSource]):
@@ -230,7 +240,7 @@ def _load_data_signals(sources: List[DataSource]):
 data_signals_by_key = {d.key: d for d in data_signals}
 # also add the resolved signal version to the signal lookup
 for d in data_signals:
-    source = data_source_by_id.get(d.source)
+    source = data_sources_by_id.get(d.source)
     if source and source.uses_db_alias:
         data_signals_by_key[(source.db_source, d.signal)] = d
 
@@ -265,7 +275,7 @@ def create_source_signal_alias_mapper(source_signals: List[SourceSignalPair]) ->
     alias_to_data_sources: Dict[str, List[DataSource]] = {}
     transformed_pairs: List[SourceSignalPair] = []
     for pair in source_signals:
-        source = data_source_by_id.get(pair.source)
+        source = data_sources_by_id.get(pair.source)
         if not source or not source.uses_db_alias:
             transformed_pairs.append(pair)
             continue
@@ -299,3 +309,159 @@ def map_row(source: str, signal: str) -> str:
         return signal_source.source
 
     return transformed_pairs, map_row
+
+
+def _resolve_all_signals(source_signals: Union[SourceSignalPair, List[SourceSignalPair]], data_sources_by_id: Dict[str, DataSource]) -> Union[SourceSignalPair, List[SourceSignalPair]]:
+    """Expand a request for all signals to an explicit list of signal names.
+
+    Example: SourceSignalPair("jhu-csse", signal=True) would return SourceSignalPair("jhu-csse", [<list of all JHU signals>]).
+    """
+    return [SourceSignalPair("src", ["sig1", "sig2"])]
+
+
+
+def _reindex_iterable(iterable: Iterable[Dict], time_pairs: List[TimePair], fill_value: Optional[int] = None) -> Iterable:
+    """Produces an iterable that fills in gaps in the time window of another iterable.
+
+    Used to produce an iterable with a contiguous time index for time series operations.
+
+    We iterate over contiguous range of days made from time_pairs. If `iterable`, which is assumed to be sorted by its "time_value" key,
+    is missing a time_value in the range, a dummy row entry is returned with the correct date and the value fields set appropriately.
+    """
+    if time_pairs is None:
+        return iterable
+
+    day_range_index = get_day_range(time_pairs)
+    for day in day_range_index.time_values:
+        if day in next_value(iterable):
+            return next_value(iterable)
+        else:
+            yield updated_default_value
+
+
+def _get_base_signal_transform(signal: Union[DataSignal, Tuple[str, str]], data_signals_by_key: Dict[Tuple[str, str], DataSignal] = data_signals_by_key) -> Callable:
+    """Given a DataSignal, return the transformation that needs to be applied to its base signal to derive the signal."""
+    return IDENTITY
+
+
+def get_transform_types(source_signal_pairs: List[SourceSignalPair], data_sources_by_id: Dict[str, DataSource] = data_sources_by_id, data_signals_by_key: Dict[Tuple[str, str], DataSignal] = data_signals_by_key) -> Set[Callable]:
+    """Return a collection of the unique transforms required for transforming a given source-signal pair list.
+
+    Example:
+    SourceSignalPair("src", ["sig", "sig_smoothed", "sig_diff"]) would return {IDENTITY, SMOOTH, DIFF}.
+
+    Used to pad the user DB query with extra days.
+    """
+    return set([IDENTITY])
+
+
+def get_pad_length(source_signal_pairs: List[SourceSignalPair], smoother_window_length: int, data_sources_by_id: Dict[str, DataSource] = data_sources_by_id, data_signals_by_key: Dict[Tuple[str, str], DataSignal] = data_signals_by_key):
+    """Returns the size of the extra date padding needed, depending on the transformations the source-signal pair list requires.
+
+    Example:
+    If smoothing is required, we fetch an extra 6 days. If both diffing and smoothing is required on the same signal, then we fetch 7 extra days.
+
+    Used to pad the user DB query with extra days.
+    """
+    transform_types = get_transform_types(source_signal_pairs, data_sources_by_id=data_sources_by_id, data_signals_by_key=data_signals_by_key)
+    if SMOOTH in transform_types:
+        return 7
+    else:
+        return 0
+
+
+def pad_time_pairs(time_pairs: List[TimePair], pad_length: int) -> List[TimePair]:
+    """Pads a list of TimePairs with another TimePair that extends the smallest time value by the pad_length, if needed.
+
+    Example:
+    [TimePair("day", [20210407])] with pad_length 6 would return [TimePair("day", [20210407]), TimePair("day", [(20210401, 20210407)])].
+
+    Used to pad the user DB query with extra days.
+    """
+    return [TimePair("day", [(20210401, 20210407)])]
+
+
+def pad_time_window(time_window: Tuple[int, int], pad_length: int) -> Tuple[int, int]:
+    """Extend a time window on the left by pad_length.
+
+    Example:
+    (20210407, 20210413) with pad_length 6 would return (20210401, 20210413).
+
+    Used to pad the user DB query with extra days.
+    """
+    return (20210401, 20210407)
+
+
+def get_day_range(time_pairs: Union[TimePair, List[TimePair]]) -> TimePair:
+    """Combine a list of TimePairs into a single contiguous, explicit TimePair object.
+
+    Example:
+    [TimePair("day", [20210407, 20210408]), TimePair("day", [(20210408, 20210410)])] would return
+    TimePair("day", [20210407, 20210408, 20210409, 20210410]).
+
+    Used to produce a contiguous time index for time series operations.
+    """
+    return TimePair("day", [20210407, 20210408, 20210409, 20210410])
+
+
+def _generate_transformed_rows(
+    parsed_rows: Iterable[Dict], time_pairs: Optional[List[TimePair]] = None, transform_dict: Optional[Dict[str, List[Tuple[str, str]]]]=None, transform_args: Optional[Dict] = None, group_keyfunc: Optional[Callable] = None, data_signals_by_key: Dict[Tuple[str, str], DataSignal] = data_signals_by_key,
+) -> Iterable[Dict]:
+    """Applies time-series transformations to streamed rows from a database.
+
+    Parameters:
+    parsed_rows: Iterable[Dict]
+        Streamed rows from the database.
+    time_pairs: Optional[List[TimePair]], default None
+        A list of TimePairs, which can be used to create a continguous time index for time-series operations.
+        The min and max dates in the TimePairs list is used.
+    transform_dict: Optional[Dict[Tuple[str, str], List[Tuple[str, str]]]], default None
+        A dictionary mapping base sources to a list of their derived signals that the user wishes to query.
+        For example, transform_dict may be {("jhu-csse", "confirmed_cumulative_num): [("jhu-csse", "confirmed_incidence_num"), ("jhu-csse", "confirmed_7dav_incidence_num")]}.
+    transform_args: Optional[Dict], default None
+        A dictionary of keyword arguments for the transformer functions.
+    group_keyfunc: Optional[Callable], default None
+        The groupby function to use to order the streamed rows. Note that Python groupby does not do any sorting, so
+        parsed_rows are assumed to be sorted in accord with this groupby.
+    data_signals_by_key: Dict[Tuple[str, str], DataSignal], default data_signals_by_key
+        The dictionary of DataSignals which is used to find the base signal transforms
+
+    Yields:
+    transformed rows: Dict
+        The transformed rows returned in an interleaved fashion. Non-transformed rows have the IDENTITY operation applied.
+    """
+    if not transform_args:
+        transform_args = dict()
+    if not transform_dict:
+        transform_dict = dict()
+    if not group_keyfunc:
+        group_keyfunc = lambda row: (row["geo_type"], row["geo_value"], row["source"], row["signal"])
+
+    try:
+        for key, group in groupby(parsed_rows, group_keyfunc):
+            _, _, source_name, signal_name = key
+            # Extract the list of derived signals.
+            # Create a list of source-signal pairs along with the transformation required for the signal.
+            # Put the current time series on a contiguous time index.
+            # Create copies of the iterable, with smart memory usage.
+            # Create a list of transformed group iterables, remembering their derived name as needed.
+            # Traverse through the transformed iterables in an interleaved fashion, which makes sure that only a small window
+            # of the original iterable (group) is stored in memory.
+            for row in transform_group(group):
+                yield row
+    except Exception as e:
+        print(f"Tranformation encountered error of type {type(e)}, with message {e}. Yielding None and stopping.")
+        yield None
+
+
+def get_basename_signals(source_signal_pairs: List[SourceSignalPair], data_sources_by_id: Dict[str, DataSource] = data_sources_by_id, data_signals_by_key: Dict[Tuple[str, str], DataSignal] = data_signals_by_key) -> Tuple[List[SourceSignalPair], Generator]:
+    """From a list of SourceSignalPairs, return the base signals required to derive them and a transformation function to take a stream
+    of the base signals and return the transformed signals.
+
+    Example:
+    SourceSignalPair("src", signal=["sig_base", "sig_smoothed"]) would return SourceSignalPair("src", signal=["sig_base"]) and a transformation function
+    that will take the returned database query for "sig_base" and return both the base time series and the smoothed time series.
+    """
+    transform_dict = {("src", "sig_base"): [("src", "sig_base"), ("src", "sig_smoothed")]}
+    row_transform_generator = partial(_generate_transformed_rows, transform_dict=transform_dict, data_signals_by_key=data_signals_by_key)
+    return SourceSignalPair("src", signal=["sig_base"]), row_transform_generator