FPM-566 support alt_df in AltData infilling

Author: dolegi

src/time_stream/infill.py

@@ -364,37 +364,39 @@ class AltData(InfillMethod):
 
     name = "alt_data"
 
-    def __init__(self, alt_data_column: str, correction_factor: float = 1.0):
+    def __init__(self, alt_data_column: str, correction_factor: float = 1.0, alt_df: pl.DataFrame | None = None):
         """Initialize the alternative data infill method.
 
         Args:
             alt_data_column: The name of the column providing the alternative data.
             correction_factor: An optional correction factor to apply to the alternative data.
+            alt_df: The DataFrame containing the alternative data.
         """
         self.alt_data_column = alt_data_column
         self.correction_factor = correction_factor
+        self.alt_df = alt_df
 
-    def _fill(self, df: pl.DataFrame, infill_column: str, alt_df: Optional[pl.DataFrame]) -> pl.DataFrame:
+    def _fill(self, df: pl.DataFrame, infill_column: str) -> pl.DataFrame:
         """Fill missing values using data from the alternative column.
 
         Args:
             df: The DataFrame to infill.
             infill_column: The column to infill.
-            alt_df: The DataFrame containing the alternative data.
 
         Returns:
             pl.DataFrame with infilled values.
         """
-        if alt_df is None:
+        breakpoint()
+        if self.alt_df is None:
             check_columns_in_dataframe(df, [self.alt_data_column])
         else:
-            check_columns_in_dataframe(alt_df, ['time', self.alt_data_column])
+            check_columns_in_dataframe(self.alt_df, ['time', self.alt_data_column])
 
             if self.alt_data_column in df.columns:
                 raise ValueError(f"Column {self.alt_data_column} already exists in the main dataframe.")
 
             df = df.join(
-                alt_df.select(['time', self.alt_data_column]),
+                self.alt_df.select(['time', self.alt_data_column]),
                 on='time',
                 how="left",
                 suffix="_alt"
@@ -407,7 +409,7 @@ def _fill(self, df: pl.DataFrame, infill_column: str, alt_df: Optional[pl.DataFr
             .alias(self._infilled_column_name(infill_column))
         )
 
-        if alt_df is not None:
+        if self.alt_df is not None:
             infilled = infilled.drop(self.alt_data_column)
 
         return infilled

tests/time_stream/test_infill.py

@@ -492,3 +492,44 @@ def test_alt_data_infill_restricting_date_range(self) -> None:
         )
         expected_df = self.df.with_columns(pl.Series("values", [1.0, 20.0, 3.0, None, 5.0]))
         assert_frame_equal(result_df, expected_df, check_column_order=False)
+
+    def test_alt_data_infill_with_alt_data_provided(self) -> None:
+        """Test infilling from a provided alternative DataFrame."""
+        alt_df = pl.DataFrame(
+            {
+                "timestamp": self.df["timestamp"],
+                "alt_values_df": [11.0, 22.0, 33.0, 44.0, 55.0],
+            }
+        )
+        infiller = AltData(alt_data_column="alt_values_df", alt_df=alt_df)
+        result_df = infiller.apply(self.tf.df, self.tf.time_name, self.tf.periodicity, "values")
+        expected_df = self.df.with_columns(pl.Series("values", [1.0, 22.0, 3.0, 44.0, 5.0]))
+        assert_frame_equal(result_df, expected_df, check_column_order=False)
+
+    def test_alt_data_infill_with_alt_data_missing_time_column(self) -> None:
+        """Test error when provided alt_data is missing the time column."""
+        alt_df = pl.DataFrame({"alt_values_df": [11.0, 22.0, 33.0, 44.0, 55.0]})
+        infiller = AltData(alt_data_column="alt_values", alt_df=alt_df)
+        with self.assertRaises(ColumnNotFoundError):
+            infiller.apply(self.tf.df, self.tf.time_name, self.tf.periodicity, "values")
+
+    def test_alt_data_infill_with_alt_data_missing_data_column(self) -> None:
+        """Test error when provided alt_data is missing the data column."""
+        alt_df = pl.DataFrame({"time": self.df["timestamp"]})
+        infiller = AltData(alt_data_column="non_existent_column", alt_df=alt_df)
+        with self.assertRaises(ColumnNotFoundError):
+            infiller.apply(self.tf.df, self.tf.time_name, self.tf.periodicity, "values")
+
+    def test_alt_data_infill_with_alt_data_and_column_in_main_df(self) -> None:
+        """Test that alt_data is prioritized when column name exists in main df."""
+        alt_df = pl.DataFrame(
+            {
+                "timestamp": self.df["timestamp"],
+                "alt_values": [11.0, 22.0, 33.0, 44.0, 55.0],
+            }
+        )
+        infiller = AltData(alt_data_column="alt_values", alt_df=alt_df)
+
+        with self.assertRaises(ValueError):
+            infiller.apply(self.tf.df, self.tf.time_name, self.tf.periodicity, "values")
+