Function normalize_timestamp_columns

sdk/python/feast/infra/ray_shared_utils.py:213–249

(
    data: Union[pd.DataFrame, Dataset, Any],
    columns: Union[str, List[str]],
    inplace: bool = False,
    exclude_columns: Optional[List[str]] = None,
)

def normalize_timestamp_columns(
    data: Union[pd.DataFrame, Dataset, Any],
    columns: Union[str, List[str]],
    inplace: bool = False,
    exclude_columns: Optional[List[str]] = None,
) -> Union[pd.DataFrame, Dataset, Any]:
    column_list = [columns] if isinstance(columns, str) else columns
    exclude_columns = exclude_columns or []

    def apply_normalization(series: pd.Series) -> pd.Series:
        return (
            pd.to_datetime(series, utc=True, errors="coerce")
            .dt.floor("s")
            .astype("datetime64[ns, UTC]")
        )

    if is_ray_data(data):

        def normalize_batch(batch: pd.DataFrame) -> pd.DataFrame:
            for column in column_list:
                if (
                    not batch.empty
                    and column in batch.columns
                    and column not in exclude_columns
                ):
                    batch[column] = apply_normalization(batch[column])
            return batch

        return data.map_batches(normalize_batch, batch_format="pandas")
    else:
        assert isinstance(data, pd.DataFrame)
        if not inplace:
            data = data.copy()
        for column in column_list:
            if column in data.columns and column not in exclude_columns:
                data[column] = apply_normalization(data[column])
        return data


def ensure_timestamp_compatibility(