JIT: update Pandas to build a single DataFrame per query

Author: dshemetov

src/server/endpoints/covidcast_utils/model.py

@@ -491,7 +491,6 @@ def _generate_transformed_rows(
     parsed_rows: Iterator[Dict],
     transform_dict: Optional[SignalTransforms] = None,
     transform_args: Optional[Dict] = None,
-    group_keyfunc: Optional[Callable] = None,
 ) -> Iterator[Dict]:
     """Applies time-series transformations to streamed rows from a database.
 
@@ -503,9 +502,6 @@ def _generate_transformed_rows(
         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.
 
     Yields:
     transformed rows: Dict
@@ -515,59 +511,57 @@ def _generate_transformed_rows(
         transform_args = dict()
     if not transform_dict:
         transform_dict = dict()
-    if not group_keyfunc:
-        group_keyfunc = lambda row: (row["source"], row["signal"], row["geo_type"], row["geo_value"])
 
-    for key, source_signal_geo_rows in groupby(parsed_rows, group_keyfunc):
+    # 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(parsed_rows, columns=columns)
+    for key, group_df in df.groupby(["source", "signal", "geo_type", "geo_value"]):
         base_source_name, base_signal_name, _, _ = key
         # 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[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")
+                yield from group_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))
+            
+            derived_df = group_df.set_index(["time_value"])
+            derived_df = derived_df.reindex(iterate_over_range(derived_df.index.min(), derived_df.index.max(), inclusive=True))
 
             if transform == DIFF:
-                df2["value"] = df2["value"].diff()
+                derived_df["value"] = derived_df["value"].diff()
                 window_length = 2
             elif transform == SMOOTH:
-                df2["value"] = df2["value"].rolling(7).mean()
+                derived_df["value"] = derived_df["value"].rolling(7).mean()
                 window_length = 7
             elif transform == DIFF_SMOOTH:
-                df2["value"] = df2["value"].diff().rolling(7).mean()
+                derived_df["value"] = derived_df["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"),
+            derived_df = derived_df.assign(
+                geo_type = derived_df["geo_type"].fillna(method="ffill"),
+                geo_value = derived_df["geo_value"].fillna(method="ffill"),
+                source = derived_df["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),
+                time_type = derived_df["time_type"].fillna(method="ffill"),
+                direction = derived_df["direction"].fillna(method="ffill"),
+                issue = derived_df["issue"].rolling(window_length).max(),
+                lag = derived_df["lag"].rolling(window_length).max(),
+                missing_value=np.where(derived_df["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"):
+            derived_df = derived_df.iloc[window_length - 1:]
+            for row in derived_df.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"]