[FSTORE-1424] Feature logging for spark (logicalclocks#242)

* fix untransformed features

fix feature type

use transformed features

feature logging spark

* run materialzation job

* fix

* refactor feature logging property

* fix feature logging getter

* fix feature logging getter

* fix fv loggingEnabled

* fix fv log

* ruff fix

* fix test

* fix parse_schema_feature_group

* fix copy polars df

* fix copy polars df

* address comments

* style

Author: kennethmhc

python/hsfs/core/feature_group_engine.py

@@ -128,7 +128,7 @@ def insert(
         validation_options: dict = None,
     ):
         dataframe_features = engine.get_instance().parse_schema_feature_group(
-            feature_dataframe, feature_group.time_travel_format
+            feature_dataframe, feature_group.time_travel_format, features=feature_group.features
         )
         dataframe_features = (
             self._update_feature_group_schema_on_demand_transformations(

python/hsfs/core/feature_logging.py

@@ -21,8 +21,8 @@ def from_response_json(cls, json_dict: Dict[str, Any]) -> "FeatureLogging":
         from hsfs.feature_group import FeatureGroup  # avoid circular import
 
         json_decamelized = humps.decamelize(json_dict)
-        transformed_features = json_decamelized.get("transformed_log")
-        untransformed_features = json_decamelized.get("untransformed_log")
+        transformed_features = json_decamelized.get("transformed_log_fg")
+        untransformed_features = json_decamelized.get("untransformed_log_fg")
         if transformed_features:
             transformed_features = FeatureGroup.from_response_json(transformed_features)
         if untransformed_features:
@@ -33,6 +33,11 @@ def from_response_json(cls, json_dict: Dict[str, Any]) -> "FeatureLogging":
             json_decamelized.get("id"), transformed_features, untransformed_features
         )
 
+    def update(self, others):
+        self._transformed_features = others.transformed_features
+        self._untransformed_features = others.untransformed_features
+        return self
+
     @property
     def transformed_features(self) -> "feature_group.FeatureGroup":
         return self._transformed_features
@@ -41,15 +46,21 @@ def transformed_features(self) -> "feature_group.FeatureGroup":
     def untransformed_features(self) -> "feature_group.FeatureGroup":
         return self._untransformed_features
 
+    def get_feature_group(self, transformed):
+        if transformed:
+            return self._transformed_features
+        else:
+            return self._untransformed_features
+
     @property
     def id(self) -> str:
         return self._id
 
     def to_dict(self):
         return {
             "id": self._id,
-            "transformed_log": self._transformed_features,
-            "untransformed_log": self._untransformed_features,
+            "transformed_log_fg": self._transformed_features,
+            "untransformed_log_fg": self._untransformed_features,
         }
 
     def json(self) -> Dict[str, Any]:

python/hsfs/core/feature_view_engine.py

@@ -17,8 +17,10 @@
 
 import datetime
 import warnings
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, TypeVar, Union
 
+import numpy as np
+import pandas as pd
 from hsfs import (
     client,
     engine,
@@ -40,6 +42,7 @@
     training_dataset_engine,
 )
 from hsfs.core.feature_logging import FeatureLogging
+from hsfs.feature_view import FeatureView
 from hsfs.training_dataset_split import TrainingDatasetSplit
 
 
@@ -967,39 +970,46 @@ def get_feature_logging(self, fv):
         )
 
     def _get_logging_fg(self, fv, transformed):
-        feature_logging = self.get_feature_logging(fv)
-        if transformed:
-            return feature_logging.transformed_features
-        else:
-            return feature_logging.untransformed_features
+        return self.get_feature_logging(fv).get_feature_group(transformed)
 
     def log_features(
         self,
-        fv,
-        features,
-        prediction=None,
-        transformed=False,
-        write_options=None,
-        training_dataset_version=None,
+        fv: FeatureView,
+        feature_logging: FeatureLogging,
+        features_rows: Union[
+            pd.DataFrame, list[list], np.ndarray, TypeVar("pyspark.sql.DataFrame")
+        ],
+        predictions: Optional[Union[pd.DataFrame, list[list], np.ndarray]] = None,
+        transformed: Optional[bool] = False,
+        write_options: Optional[Dict[str, Any]] = None,
+        training_dataset_version: Optional[int] = None,
         hsml_model=None,
     ):
         default_write_options = {
             "start_offline_materialization": False,
         }
         if write_options:
             default_write_options.update(write_options)
-        fg = self._get_logging_fg(fv, transformed)
+        fg = feature_logging.get_feature_group(transformed)
+        td_predictions = [feature for feature in fv.features if feature.label]
+        td_predictions_names = set([feature.name for feature in td_predictions])
+        if transformed:
+            td_features = [feature_name for feature_name in fv.transformed_features if feature_name not in td_predictions_names]
+        else:
+            td_features = [feature.name for feature in
+                           fv.features if
+                           feature.name not in td_predictions_names]
         df = engine.get_instance().get_feature_logging_df(
-            fg,
-            features,
-            [feature for feature in fv.features if not feature.label],
-            [feature for feature in fv.features if feature.label],
-            FeatureViewEngine._LOG_TD_VERSION,
-            FeatureViewEngine._LOG_TIME,
-            FeatureViewEngine._HSML_MODEL,
-            prediction,
-            training_dataset_version,
-            hsml_model,
+            features_rows,
+            fg=fg,
+            td_features=td_features,
+            td_predictions=td_predictions,
+            td_col_name=FeatureViewEngine._LOG_TD_VERSION,
+            time_col_name=FeatureViewEngine._LOG_TIME,
+            model_col_name=FeatureViewEngine._HSML_MODEL,
+            predictions=predictions,
+            training_dataset_version=training_dataset_version,
+            hsml_model=hsml_model,
         )
         return fg.insert(df, write_options=default_write_options)
 
@@ -1038,9 +1048,7 @@ def read_feature_logs(
             query = query.filter(
                 self._convert_to_log_fg_filter(fg, fv, filter, fv_feat_name_map)
             )
-        df = query.read()
-        df = df.drop(["log_id", FeatureViewEngine._LOG_TIME], axis=1)
-        return df
+        return engine.get_instance().read_feature_log(query)
 
     @staticmethod
     def get_hsml_model_value(hsml_model):
@@ -1099,15 +1107,19 @@ def pause_logging(self, fv):
     def resume_logging(self, fv):
         self._feature_view_api.resume_feature_logging(fv.name, fv.version)
 
-    def materialize_feature_logs(self, fv, wait):
-        jobs = self._feature_view_api.materialize_feature_logging(fv.name, fv.version)
+    def materialize_feature_logs(self, fv, wait, transform):
+        if transform is None:
+            jobs = [self._get_logging_fg(fv, True).materialization_job,
+                    self._get_logging_fg(fv, False).materialization_job]
+        else:
+            jobs = [self._get_logging_fg(fv, transform).materialization_job]
+        for job in jobs:
+            job.run(await_termination=False)
         if wait:
             for job in jobs:
-                try:
-                    job._wait_for_job(wait)
-                except Exception:
-                    pass
+                job._wait_for_job(wait)
         return jobs
 
-    def delete_feature_logs(self, fv, transformed):
+    def delete_feature_logs(self, fv, feature_logging, transformed):
         self._feature_view_api.delete_feature_logs(fv.name, fv.version, transformed)
+        feature_logging.update(self.get_feature_logging(fv))

python/hsfs/engine/python.py

@@ -734,20 +734,28 @@ def parse_schema_feature_group(
         self,
         dataframe: Union[pd.DataFrame, pl.DataFrame],
         time_travel_format: Optional[str] = None,
+        features: Optional[List[feature.Feature]] = None,
     ) -> List[feature.Feature]:
+        feature_type_map = {}
+        if features:
+            for _feature in features:
+                feature_type_map[_feature.name] = _feature.type
         if isinstance(dataframe, pd.DataFrame):
             arrow_schema = pa.Schema.from_pandas(dataframe, preserve_index=False)
         elif isinstance(dataframe, pl.DataFrame) or isinstance(
             dataframe, pl.dataframe.frame.DataFrame
         ):
             arrow_schema = dataframe.to_arrow().schema
         features = []
-        for feat_name in arrow_schema.names:
+        for i in range(len(arrow_schema.names)):
+            feat_name = arrow_schema.names[i]
             name = util.autofix_feature_name(feat_name)
             try:
-                converted_type = convert_pandas_dtype_to_offline_type(
-                    arrow_schema.field(feat_name).type
-                )
+                pd_type = arrow_schema.field(feat_name).type
+                if pa.types.is_null(pd_type) and feature_type_map.get(name):
+                    converted_type = feature_type_map.get(name)
+                else:
+                    converted_type = convert_pandas_dtype_to_offline_type(pd_type)
             except ValueError as e:
                 raise FeatureStoreException(f"Feature '{name}': {str(e)}") from e
             features.append(feature.Feature(name, converted_type))
@@ -1422,7 +1430,7 @@ def _start_offline_materialization(offline_write_options: Dict[str, Any]) -> boo
             return True
 
     @staticmethod
-    def _convert_feature_log_to_df(feature_log, cols):
+    def _convert_feature_log_to_df(feature_log, cols) -> pd.DataFrame:
         if feature_log is None and cols:
             return pd.DataFrame(columns=cols)
         if not (
@@ -1442,40 +1450,40 @@ def _convert_feature_log_to_df(feature_log, cols):
 
             return pd.DataFrame(feature_log, columns=cols)
         else:
-            return feature_log.copy(deep=False)
+            if isinstance(feature_log, pl.DataFrame):
+                return feature_log.clone().to_pandas()
+            elif isinstance(feature_log, pd.DataFrame):
+                return feature_log.copy(deep=False)
 
     @staticmethod
     def get_feature_logging_df(
-        fg,
-        features,
-        fg_features: List[TrainingDatasetFeature],
-        td_predictions: List[TrainingDatasetFeature],
-        td_col_name,
-        time_col_name,
-        model_col_name,
-        prediction=None,
-        training_dataset_version=None,
+        features: Union[pd.DataFrame, list[list], np.ndarray],
+        fg: FeatureGroup = None,
+        td_features: List[str] = None,
+        td_predictions: List[TrainingDatasetFeature] = None,
+        td_col_name: Optional[str] = None,
+        time_col_name: Optional[str] = None,
+        model_col_name: Optional[str] = None,
+        predictions: Optional[Union[pd.DataFrame, list[list], np.ndarray]] = None,
+        training_dataset_version: Optional[int] = None,
         hsml_model=None,
     ) -> pd.DataFrame:
-        import uuid
-
         features = Engine._convert_feature_log_to_df(
-            features, [f.name for f in fg_features]
+            features, td_features
         )
         if td_predictions:
-            prediction = Engine._convert_feature_log_to_df(
-                prediction, [f.name for f in td_predictions]
+            predictions = Engine._convert_feature_log_to_df(
+                predictions, [f.name for f in td_predictions]
             )
             for f in td_predictions:
-                prediction[f.name] = Engine._cast_column_to_offline_type(
-                    prediction[f.name], f.type
+                predictions[f.name] = cast_column_to_offline_type(
+                    predictions[f.name], f.type
                 )
-            if not set(prediction.columns).intersection(set(features.columns)):
-                features = pd.concat([features, prediction], axis=1)
-        # need to case the column type as if it is None, type cannot be inferred.
-        features[td_col_name] = Engine._cast_column_to_offline_type(
-            pd.Series([training_dataset_version for _ in range(len(features))]),
-            fg.get_feature(td_col_name).type,
+            if not set(predictions.columns).intersection(set(features.columns)):
+                features = pd.concat([features, predictions], axis=1)
+
+        features[td_col_name] = pd.Series(
+            [training_dataset_version for _ in range(len(features))]
         )
         # _cast_column_to_offline_type cannot cast string type
         features[model_col_name] = pd.Series(
@@ -1488,9 +1496,12 @@ def get_feature_logging_df(
             dtype=pd.StringDtype(),
         )
         now = datetime.now()
-        features[time_col_name] = Engine._cast_column_to_offline_type(
-            pd.Series([now for _ in range(len(features))]),
-            fg.get_feature(time_col_name).type,
-        )
+
+        features[time_col_name] = pd.Series([now for _ in range(len(features))])
         features["log_id"] = [str(uuid.uuid4()) for _ in range(len(features))]
         return features[[feat.name for feat in fg.features]]
+
+    @staticmethod
+    def read_feature_log(query):
+        df = query.read()
+        return df.drop(["log_id", FeatureViewEngine._LOG_TIME], axis=1)

python/hsfs/engine/spark.py

@@ -21,10 +21,14 @@
 import os
 import re
 import shutil
+import uuid
 import warnings
 from datetime import date, datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypeVar, Union
 
+from hsfs.core.feature_view_engine import FeatureViewEngine
+from hsfs.training_dataset_feature import TrainingDatasetFeature
+
 
 if TYPE_CHECKING:
     import great_expectations
@@ -1150,6 +1154,7 @@ def parse_schema_feature_group(
         self,
         dataframe,
         time_travel_format=None,
+        **kwargs,
     ):
         features = []
 
@@ -1477,11 +1482,55 @@ def cast_columns(df, schema, online=False):
     def is_connector_type_supported(type):
         return True
 
-    @staticmethod
-    def get_feature_logging_df(features, prediction=None):
+    def get_feature_logging_df(
+        self,
+        features: Union[
+            pd.DataFrame, list[list], np.ndarray, TypeVar("pyspark.sql.DataFrame")
+        ],
+        fg: fg_mod.FeatureGroup = None,
+        td_features: List[str] = None,
+        td_predictions: List[TrainingDatasetFeature] = None,
+        td_col_name: Optional[str] = None,
+        time_col_name: Optional[str] = None,
+        model_col_name: Optional[str] = None,
+        predictions: Optional[Union[pd.DataFrame, list[list], np.ndarray]] = None,
+        training_dataset_version: Optional[int] = None,
+        hsml_model=None,
+        **kwargs,
+    ):
         # do not take prediction separately because spark ml framework usually return feature together with the prediction
         # and it is costly to join them back
-        return features
+        df = self.convert_to_default_dataframe(features)
+        if td_predictions:
+            for f in td_predictions:
+                if f.name not in df.columns:
+                    df = df.withColumn(
+                        f.name,
+                        lit(None).cast(
+                            Engine._convert_offline_type_to_spark_type(f.type)
+                        ),
+                    )
+
+        uuid_udf = udf(lambda: str(uuid.uuid4()), StringType())
+
+        # Add new columns to the DataFrame
+        df = df.withColumn(td_col_name, lit(training_dataset_version).cast(LongType()))
+        if hsml_model is not None:
+            hsml_str = FeatureViewEngine.get_hsml_model_value(hsml_model)
+        else:
+            hsml_str = None
+        df = df.withColumn(model_col_name, lit(hsml_str).cast(StringType()))
+        now = datetime.now()
+        df = df.withColumn(time_col_name, lit(now).cast(TimestampType()))
+        df = df.withColumn("log_id", uuid_udf())
+
+        # Select the required columns
+        return df.select(*[feat.name for feat in fg.features])
+
+    @staticmethod
+    def read_feature_log(query):
+        df = query.read()
+        return df.drop("log_id", FeatureViewEngine._LOG_TIME)
 
 
 class SchemaError(Exception):