from feast.importer import import_class

TRANSFORMATION_CLASS_FOR_TYPE = {

"python": "feast.transformation.python_transformation.PythonTransformation",

"pandas": "feast.transformation.pandas_transformation.PandasTransformation",

"substrait": "feast.transformation.substrait_transformation.SubstraitTransformation",

"sql": "feast.transformation.sql_transformation.SQLTransformation",

"spark_sql": "feast.transformation.spark_transformation.SparkTransformation",

"spark": "feast.transformation.spark_transformation.SparkTransformation",

"ray": "feast.transformation.ray_transformation.RayTransformation",

}

def get_transformation_class_from_type(transformation_type: str):

if transformation_type in TRANSFORMATION_CLASS_FOR_TYPE:

transformation_type = TRANSFORMATION_CLASS_FOR_TYPE[transformation_type]

elif not transformation_type.endswith("Transformation"):

raise ValueError(

f"Invalid transformation type: {transformation_type}. Choose from {list(TRANSFORMATION_CLASS_FOR_TYPE.keys())}."

)

module_name, transformation_class_type = transformation_type.rsplit(".", 1)

return import_class(

module_name, transformation_class_type, transformation_class_type

)