FPM-566 use consistant time column name based on what is originally provided to time-stream

Author: dolegi

src/time_stream/infill.py

@@ -15,7 +15,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 import polars as pl
@@ -386,20 +386,20 @@ def _fill(self, df: pl.DataFrame, infill_column: str) -> pl.DataFrame:
         Returns:
             pl.DataFrame with infilled values.
         """
-        breakpoint()
         if self.alt_df is None:
             check_columns_in_dataframe(df, [self.alt_data_column])
         else:
-            check_columns_in_dataframe(self.alt_df, ['time', self.alt_data_column])
+            time_column_name = next(iter(df.schema))
+            check_columns_in_dataframe(self.alt_df, [time_column_name, 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(
-                self.alt_df.select(['time', self.alt_data_column]),
-                on='time',
+                self.alt_df.select([time_column_name, self.alt_data_column]),
+                on=time_column_name,
                 how="left",
-                suffix="_alt"
+                suffix="_alt",
             )
 
         infilled = df.with_columns(

tests/time_stream/test_infill.py

@@ -532,4 +532,3 @@ def test_alt_data_infill_with_alt_data_and_column_in_main_df(self) -> None:
 
         with self.assertRaises(ValueError):
             infiller.apply(self.tf.df, self.tf.time_name, self.tf.periodicity, "values")
-