JIT: add Pandas JIT approach

Author: dshemetov

integrations/server/test_covidcast.py

@@ -138,7 +138,7 @@ def test_csv_format(self):
       **{'format':'csv'}
     )
 
-    # TODO: This is a mess because of api.php.
+    # TODO: This is a mess because of api.php. Or maybe it's just a mess.
     column_order = [
       "geo_value", "signal", "time_value", "direction", "issue", "lag", "missing_value",
       "missing_stderr", "missing_sample_size", "value", "stderr", "sample_size"

src/server/endpoints/covidcast.py

@@ -7,7 +7,7 @@
 from flask import Blueprint, request
 from flask.json import loads, jsonify
 from more_itertools import peekable
-from numpy import nan
+from numpy import int64
 from sqlalchemy import text
 from pandas import read_csv, to_datetime
 
@@ -136,22 +136,9 @@ def parse_transform_args():
     if smoother_window_length is None:
         smoother_window_length = 7
 
-    # TODO: Add support for floats inputs here.
-    # The value to fill for missing date values.
-    pad_fill_value = extract_integer("pad_fill_value")
-    if pad_fill_value is None:
-        pad_fill_value = nan
-
-    # The value to fill for None or nan values.
-    nan_fill_value = extract_integer("nans_fill_value")
-    if nan_fill_value is None:
-        nan_fill_value = nan
-
     smoother_args = {
         "smoother_kernel": SmootherKernelValue.average,
         "smoother_window_length": smoother_window_length if isinstance(smoother_window_length, Number) and smoother_window_length <= MAX_SMOOTHER_WINDOW else MAX_SMOOTHER_WINDOW,
-        "pad_fill_value": pad_fill_value if isinstance(pad_fill_value, Number) else nan,
-        "nans_fill_value": nan_fill_value if isinstance(nan_fill_value, Number) else nan
     }
     return smoother_args
 
@@ -210,7 +197,7 @@ def alias_row(row):
 
         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)
+            transformed_rows = row_transform_generator(parsed_rows=parsed_rows, transform_args=transform_args)
             for row in transformed_rows:
                 yield alias_row(row)
     else:
@@ -300,7 +287,7 @@ def gen_trend(rows):
 
         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=[TimePair("day", [time_window])], transform_args=transform_args)
+            transformed_rows = row_transform_generator(parsed_rows=parsed_rows, transform_args=transform_args)
             for row in transformed_rows:
                 yield row
     else:
@@ -380,7 +367,7 @@ def gen_trend(rows):
 
         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=[TimePair("day", [time_window])], transform_args=transform_args)
+            transformed_rows = row_transform_generator(parsed_rows=parsed_rows, transform_args=transform_args)
             for row in transformed_rows:
                 yield row
     else:
@@ -527,7 +514,7 @@ def handle_export():
 
         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=[TimePair("day", [time_window])], transform_args=transform_args)
+            transformed_rows = row_transform_generator(parsed_rows=parsed_rows, transform_args=transform_args)
             for row in transformed_rows:
                 yield row
     else:
@@ -556,7 +543,7 @@ def parse_csv_row(i, row):
             "geo_value": row["geo_value"],
             "signal": row["signal"],
             "time_value": time_value_to_iso(row["time_value"]) if is_day else row["time_value"],
-            "issue": time_value_to_iso(row["issue"]) if is_day else row["issue"],
+            "issue": time_value_to_iso(row["issue"]) if is_day and isinstance(row["issue"], (int, int64)) else row["issue"],
             "lag": row["lag"],
             "value": row["value"],
             "stderr": row["stderr"],

src/server/endpoints/covidcast_utils/model.py

@@ -1,20 +1,20 @@
 from dataclasses import asdict, dataclass, field
 from enum import Enum
 from functools import partial
-from itertools import groupby, repeat, tee
+from itertools import groupby
 from numbers import Number
 from typing import Callable, Generator, Iterator, Optional, Dict, List, Set, Tuple, Union
 
 from pathlib import Path
 import re
-from more_itertools import flatten, interleave_longest, peekable
+from more_itertools import flatten, peekable
 import pandas as pd
 import numpy as np
 
 from delphi_utils.nancodes import Nans
 from ..._params import SourceSignalPair, TimePair
 from .smooth_diff import generate_smoothed_rows, generate_diffed_rows
-from ...utils import shift_time_value, iterate_over_ints_and_ranges
+from ...utils import shift_time_value, iterate_over_ints_and_ranges, iterate_over_range
 
 
 IDENTITY: Callable = lambda rows, **kwargs: rows
@@ -489,7 +489,6 @@ def get_day_range(time_pairs: List[TimePair]) -> Iterator[int]:
 
 def _generate_transformed_rows(
     parsed_rows: Iterator[Dict],
-    time_pairs: Optional[List[TimePair]] = None,
     transform_dict: Optional[SignalTransforms] = None,
     transform_args: Optional[Dict] = None,
     group_keyfunc: Optional[Callable] = None,
@@ -499,9 +498,6 @@ def _generate_transformed_rows(
     Parameters:
     parsed_rows: Iterator[Dict]
         An iterator streaming rows from a database query. Assumed to be sorted by source, signal, geo_type, geo_value, time_type, and time_value.
-    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[SignalTransforms], 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")]}.
@@ -527,18 +523,56 @@ def _generate_transformed_rows(
         # Extract the list of derived signals; if a signal is not in the dictionary, then use the identity map.
         derived_signal_transform_map: SourceSignalPair = transform_dict.get(SourceSignalPair(base_source_name, [base_signal_name]), SourceSignalPair(base_source_name, [base_signal_name]))
         # Create a list of source-signal pairs along with the transformation required for the signal.
-        signal_names_and_transforms: List[Tuple[Tuple[str, str], Callable]] = [(derived_signal, _get_base_signal_transform((base_source_name, derived_signal))) for derived_signal in derived_signal_transform_map.signal]
-        # Put the current time series on a contiguous time index.
-        source_signal_geo_rows = _reindex_iterable(source_signal_geo_rows, time_pairs, fill_value=transform_args.get("pad_fill_value"))
-        # Create copies of the iterable, with smart memory usage.
-        source_signal_geo_rows_copies: Iterator[Iterator[Dict]] = tee(source_signal_geo_rows, len(signal_names_and_transforms))
-        # Create a list of transformed group iterables, remembering their derived name as needed.
-        transformed_signals_iterator: Iterator[Tuple[str, Iterator[Dict]]] = (zip(repeat(derived_signal), transform(rows, **transform_args)) for (derived_signal, transform), rows in zip(signal_names_and_transforms, source_signal_geo_rows_copies))
-        # 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 derived_signal_name, row in interleave_longest(*transformed_signals_iterator):
-            row["signal"] = derived_signal_name
-            yield row
+        signal_names_and_transforms: List[Tuple[str, Callable]] = [(derived_signal, _get_base_signal_transform((base_source_name, derived_signal))) for derived_signal in derived_signal_transform_map.signal]
+
+        # TODO: Fix these to come as an argument.
+        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"]
+        columns = fields_string + fields_int + fields_float
+        df = pd.DataFrame.from_records(source_signal_geo_rows, columns=columns)
+        for derived_signal, transform in signal_names_and_transforms:
+            if transform == IDENTITY:
+                yield from df.to_dict(orient="records")
+                continue
+
+            df2 = df.set_index(["time_value"])
+            df2 = df2.reindex(iterate_over_range(df2.index.min(), df2.index.max(), inclusive=True))
+
+            if transform == DIFF:
+                df2["value"] = df2["value"].diff()
+                window_length = 2
+            elif transform == SMOOTH:
+                df2["value"] = df2["value"].rolling(7).mean()
+                window_length = 7
+            elif transform == DIFF_SMOOTH:
+                df2["value"] = df2["value"].diff().rolling(7).mean()
+                window_length = 8
+            else:
+                raise ValueError(f"Unknown transform for {derived_signal}.")
+
+            df2 = df2.assign(
+                geo_type = df2["geo_type"].fillna(method="ffill"),
+                geo_value = df2["geo_value"].fillna(method="ffill"),
+                source = df2["source"].fillna(method="ffill"),
+                signal = derived_signal,
+                time_type = df2["time_type"].fillna(method="ffill"),
+                direction = df2["direction"].fillna(method="ffill"),
+                issue = df2["issue"].rolling(window_length).max(),
+                lag = df2["lag"].rolling(window_length).max(),
+                missing_value=np.where(df2["value"].isna(), Nans.NOT_APPLICABLE, Nans.NOT_MISSING),
+                missing_stderr=Nans.NOT_APPLICABLE,
+                missing_sample_size=Nans.NOT_APPLICABLE,
+                stderr=np.nan,
+                sample_size=np.nan,
+            )
+            df2 = df2.iloc[window_length - 1:]
+            for row in df2.reset_index().to_dict(orient="records"):
+                row.update({
+                    "issue": int(row["issue"]) if not np.isnan(row["issue"]) else row["issue"], 
+                    "lag": int(row["lag"]) if not np.isnan(row["lag"]) else row["lag"]
+                })
+                yield row
 
 
 def get_basename_signal_and_jit_generator(source_signal_pairs: List[SourceSignalPair], transform_args: Optional[Dict[str, Union[str, int]]] = None) -> Tuple[List[SourceSignalPair], Generator]:

tests/server/endpoints/covidcast_utils/test_model.py

@@ -6,7 +6,7 @@
 from more_itertools import interleave_longest, windowed
 from pandas.testing import assert_frame_equal
 
-from delphi.epidata.acquisition.covidcast.covidcast_row import CovidcastRows
+from delphi.epidata.acquisition.covidcast.covidcast_row import CovidcastRows, assert_frame_equal_no_order
 from delphi.epidata.server._params import SourceSignalPair, TimePair
 from delphi.epidata.server.endpoints.covidcast_utils.model import (
     DIFF,
@@ -163,7 +163,7 @@ def test__generate_transformed_rows(self):
                 missing_sample_size=[Nans.NOT_APPLICABLE] * 4,
             ).api_row_df
 
-            assert_frame_equal(df, expected_df)
+            assert_frame_equal_no_order(df, expected_df, index=["signal", "geo_value", "time_value"])
 
         with self.subTest("smoothed and diffed signals on one base test"):
             data = CovidcastRows.from_args(
@@ -184,11 +184,7 @@ def test__generate_transformed_rows(self):
                 sample_size=[None] * 13,
             ).api_row_df
 
-            # Test no order.
-            idx = ["source", "signal", "time_value"]
-            assert_frame_equal(df.set_index(idx).sort_index(), expected_df.set_index(idx).sort_index())
-            # Test order.
-            assert_frame_equal(df, expected_df)
+            assert_frame_equal_no_order(df, expected_df, index=["signal", "geo_value", "time_value"])
 
         with self.subTest("smoothed and diffed signal on two non-continguous regions"):
             data = CovidcastRows.from_args(
@@ -199,9 +195,8 @@ def test__generate_transformed_rows(self):
                 sample_size=range(15),
             ).api_row_df
             transform_dict = {SourceSignalPair("src", ["sig_base"]): SourceSignalPair("src", ["sig_diff", "sig_smooth"])}
-            time_pairs = [TimePair("day", [(20210501, 20210520)])]
             df = CovidcastRows.from_records(
-                _generate_transformed_rows(data.to_dict(orient="records"), time_pairs=time_pairs, transform_dict=transform_dict)
+                _generate_transformed_rows(data.to_dict(orient="records"), transform_dict=transform_dict)
             ).api_row_df
 
             filled_values = data.value.to_list()[:10] + [None] * 5 + data.value.to_list()[10:]
@@ -215,11 +210,8 @@ def test__generate_transformed_rows(self):
                 sample_size=[None] * 33,
                 issue=interleave_longest(_reindex_windowed(filled_time_values, 2), _reindex_windowed(filled_time_values, 7)),
             ).api_row_df
-            # Test no order.
-            idx = ["source", "signal", "time_value"]
-            assert_frame_equal(df.set_index(idx).sort_index(), expected_df.set_index(idx).sort_index())
-            # Test order.
-            assert_frame_equal(df, expected_df)
+
+            assert_frame_equal_no_order(df, expected_df, index=["signal", "geo_value", "time_value"])
         # fmt: on
 
     def test_get_basename_signals(self):
@@ -258,44 +250,46 @@ def test_get_basename_signals(self):
             ).api_row_df
             source_signal_pairs = [SourceSignalPair("src", ["sig_base", "sig_diff", "sig_other", "sig_smooth"])]
             _, row_transform_generator = get_basename_signal_and_jit_generator(source_signal_pairs)
-            time_pairs = [TimePair("day", [(20210501, 20210530)])]
-            df = CovidcastRows.from_records(row_transform_generator(data.to_dict(orient="records"), time_pairs=time_pairs)).api_row_df
+            df = CovidcastRows.from_records(row_transform_generator(data.to_dict(orient="records"))).api_row_df
 
             filled_values = list(chain(range(10), [None] * 10, range(10, 20)))
             filled_time_values = list(chain(pd.date_range("2021-05-01", "2021-05-10"), [None] * 10, pd.date_range("2021-05-21", "2021-05-30")))
 
             expected_df = CovidcastRows.from_args(
-                signal=["sig_base"] * 30 + ["sig_diff"] * 29 + ["sig_other"] * 5 + ["sig_smooth"] * 24,
+                signal=["sig_base"] * 20 + ["sig_diff"] * 29 + ["sig_other"] * 5 + ["sig_smooth"] * 24,
                 time_value=chain(
-                    pd.date_range("2021-05-01", "2021-05-30"),
+                    chain(pd.date_range("2021-05-01", "2021-05-10"), pd.date_range("2021-05-21", "2021-05-30")),
                     pd.date_range("2021-05-02", "2021-05-30"),
                     pd.date_range("2021-05-01", "2021-05-05"),
                     pd.date_range("2021-05-07", "2021-05-30")
                 ),
                 value=chain(
-                    filled_values, 
+                    range(20), 
                     _diff_rows(filled_values), 
                     range(5), 
                     _smooth_rows(filled_values)
                 ),
                 stderr=chain(
-                    chain(range(10), [None] * 10, range(10, 20)),
+                    range(20),
                     chain([None] * 29),
                     range(5),
                     chain([None] * 24),
                 ),
                 sample_size=chain(
-                    chain(range(10), [None] * 10, range(10, 20)),
+                    range(20),
                     chain([None] * 29),
                     range(5),
                     chain([None] * 24),
                 ),
-                issue=chain(filled_time_values, _reindex_windowed(filled_time_values, 2), pd.date_range("2021-05-01", "2021-05-05"), _reindex_windowed(filled_time_values, 7)),
+                issue=chain(
+                    chain(pd.date_range("2021-05-01", "2021-05-10"), pd.date_range("2021-05-21", "2021-05-30")), 
+                    _reindex_windowed(filled_time_values, 2), 
+                    pd.date_range("2021-05-01", "2021-05-05"), 
+                    _reindex_windowed(filled_time_values, 7)
+                ),
             ).api_row_df
             # fmt: on
-            # Test no order.
-            idx = ["source", "signal", "time_value"]
-            assert_frame_equal(df.set_index(idx).sort_index(), expected_df.set_index(idx).sort_index())
+            assert_frame_equal_no_order(df, expected_df, index=["signal", "geo_value", "time_value"])
 
         with self.subTest("test base, diff, smooth; multiple geos"):
             # fmt: off
@@ -341,8 +335,7 @@ def test_get_basename_signals(self):
             ).api_row_df
             # fmt: on
             # Test no order.
-            idx = ["source", "signal", "time_value"]
-            assert_frame_equal(df.set_index(idx).sort_index(), expected_df.set_index(idx).sort_index())
+            assert_frame_equal_no_order(df, expected_df, index=["signal", "geo_value", "time_value"])
 
         with self.subTest("empty iterator"):
             source_signal_pairs = [SourceSignalPair("src", ["sig_base", "sig_diff", "sig_smooth"])]