Merge pull request #85 from hubverse-org/mc/time-series-parquet-files/70

support time-series.parquet files in addition to .csv

Author: matthewcornell

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "structlog",
     "pyarrow",
     "jsonschema",
-    "hubdata>=0.1.3",
+    "hubdata>=0.2.0",
 ]
 
 [project.optional-dependencies]

requirements/requirements.txt

@@ -15,7 +15,7 @@ click==8.2.1
     #   hub-dashboard-predtimechart (pyproject.toml)
     #   hubdata
     #   pip-tools
-hubdata==0.1.3
+hubdata==0.2.0
     # via hub-dashboard-predtimechart (pyproject.toml)
 iniconfig==2.0.0
     # via pytest

src/hub_predtimechart/app/generate_target_json_files.py

@@ -51,7 +51,7 @@ def main(hub_dir, ptc_config_file, target_out_dir, regenerate):
     try:
         target_data_df = hub_config.get_target_data_df()
     except FileNotFoundError as error:
-        logger.error(f"target data file not found. {hub_config.get_target_data_file_name()=}, {error=}")
+        logger.error(f"target data file not found. {error=}")
         sys.exit(1)
 
     json_files = _generate_target_json_files(hub_config, target_data_df, target_out_dir, regenerate)
@@ -131,7 +131,7 @@ def ptc_target_data(model_task: ModelTask, target_data_df: pl.DataFrame, task_id
         if max_as_of is None:
             return None
         else:
-            target_data_df = target_data_df.filter(pl.col('as_of') == max_as_of.isoformat())
+            target_data_df = target_data_df.filter(pl.col('as_of') == max_as_of)
     else:
         # the file is one that is assumed to be updated weekly and so we can
         # assume that the effective as_of date for this file is the same as
@@ -158,30 +158,30 @@ def ptc_target_data(model_task: ModelTask, target_data_df: pl.DataFrame, task_id
     if len(target_data_df) == 0:
         return None
 
+    # date column type depends on data source: date objects from `connect_target_data()`, strings from custom CSV files.
+    # convert date objects to ISO strings for JSON serialization; pass through strings as-is.
     return {
-        'date': target_data_df[target_date_col_name].to_list(),
+        'date': [d.isoformat() if isinstance(d, date) else d for d in target_data_df[target_date_col_name].to_list()],
         'y': target_data_df[observation_col_name].to_list()
     }
 
 
-def _max_as_of_le_reference_date(target_data_df: pl.DataFrame, viz_target_id: str, reference_date: str) -> date:
+def _max_as_of_le_reference_date(target_data_df: pl.DataFrame, viz_target_id: str, reference_date: str) -> date | None:
     """
     ptc_target_data() helper
 
-    :param target_data_df: a pl.DataFrame that loaded from HubConfigPtc.target_data_file_name. assumes follows our new
+    :param target_data_df: a pl.DataFrame from connect_target_data(). assumes follows the
         time-series target data standard - has `as_of` column, etc.
     :param viz_target_id: the target of interest. via ModelTask.viz_target_id
     :param reference_date: string naming the reference_date of interest
     :return: max as_of that's <= `reference_date` for `viz_target_id`. return None if not found
     """
     reference_date = date.fromisoformat(reference_date)
-    unique_as_ofs = [date.fromisoformat(as_of) for as_of in pl.Series(target_data_df
-                                                                      .filter(pl.col('target') == viz_target_id)
-                                                                      .unique('as_of')
-                                                                      .select('as_of')
-                                                                      .sort('as_of'))]  # sort for debugging
-    le_as_ofs = [as_of for as_of in unique_as_ofs if as_of <= reference_date]
-    return max(le_as_ofs) if le_as_ofs else None
+    return (target_data_df
+            .filter(pl.col('target') == viz_target_id)
+            .filter(pl.col('as_of') <= reference_date)
+            .select(pl.col('as_of').max())
+            .item())
 
 
 #

src/hub_predtimechart/hub_config_ptc.py

@@ -8,6 +8,7 @@
 import polars as pl
 import yaml
 from hubdata import HubConnection
+from hubdata.connect_target_data import TargetType, connect_target_data
 from jsonschema import FormatChecker, ValidationError, validate
 
 from hub_predtimechart.ptc_schema import ptc_config_schema
@@ -117,10 +118,24 @@ def model_output_file_for_ref_date(self, model_id: str, reference_date: str) ->
 
     def get_target_data_df(self) -> pl.DataFrame:
         """
-        Loads the target data csv file from the hub repo for now, file path for target data is hard coded to 'target-data'.
-        Raises FileNotFoundError if target data file does not exist.
+        Loads the target data file from the hub repo. Uses `hubdata.connect_target_data()` for standard target data
+        locations (time-series.csv, time-series.parquet, or time-series/ directory). Falls back to custom file reading
+        when `target_data_file_name` is specified in the predtimechart config.
+
+        :return: target data as a polars DataFrame
+        :raises FileNotFoundError: if target data file does not exist
+        :raises ValueError: if target data file has unsupported format (custom file only)
         """
-        target_data_file_path = self.hub_path / 'target-data' / self.get_target_data_file_name()
+        # non-custom file name case: use `hubdata.connect_target_data()` for standard file locations
+        if not self.target_data_file_name:
+            try:
+                target_conn = connect_target_data(self.hub_path, TargetType.TIME_SERIES)
+                return pl.from_arrow(target_conn.to_table())
+            except RuntimeError as error:
+                raise FileNotFoundError(f"target data not found via hubdata. {error=}")
+
+        # custom file name case
+        target_data_file_path = self.hub_path / 'target-data' / self.target_data_file_name
         try:
             # the override schema handles the 'US' location (the only location that doesn't parse as Int64)
             # todo hard-coded column names
@@ -132,13 +147,6 @@ def get_target_data_df(self) -> pl.DataFrame:
             raise FileNotFoundError(f"target data file not found. {target_data_file_path=}, {error=}")
 
 
-    def get_target_data_file_name(self):
-        """
-        :return: the target data file name under the "target-data" dir to use
-        """
-        return self.target_data_file_name if self.target_data_file_name else 'time-series.csv'
-
-
 def _validate_predtimechart_config(ptc_config: dict, tasks: dict):
     """
     Validates `ptc_config` against the schema in ptc_schema.py.

tests/hub_predtimechart/test_hub_config_ptc.py

@@ -215,23 +215,6 @@ def test_task_id_text_covid19_forecast_hub():
     }
 
 
-def test_get_target_data_file_name():
-    # hub that predates the new target data standard file name. specifies file name in hub_config.target_data_file_name
-    hub_path = Path('tests/hubs/covid19-forecast-hub')
-    hub_config = HubConfigPtc(hub_path, hub_path / 'hub-config/predtimechart-config.yml')
-    assert hub_config.get_target_data_file_name() == 'covid-hospital-admissions.csv'
-
-    # hub that predates the new target data standard file name. specifies file name in hub_config.target_data_file_name
-    hub_path = Path('tests/hubs/FluSight-forecast-hub')
-    hub_config = HubConfigPtc(hub_path, hub_path / 'hub-config/predtimechart-config.yml')
-    assert hub_config.get_target_data_file_name() == 'target-hospital-admissions.csv'
-
-    # hub that uses the new target data standard file name: "target-data/time-series.csv"
-    hub_path = Path('tests/hubs/flu-metrocast')
-    hub_config = HubConfigPtc(hub_path, hub_path / 'hub-config/predtimechart-config.yml')
-    assert hub_config.get_target_data_file_name() == 'time-series.csv'
-
-
 def test_model_tasks_instance_list():
     for hub_path, exp_len in [('tests/hubs/covid19-forecast-hub', 1),
                              ('tests/hubs/example-complex-forecast-hub', 1),