98 lines
3.1 KiB
Python
98 lines
3.1 KiB
Python
import re
|
|
from functools import reduce
|
|
from typing import Union
|
|
|
|
try:
|
|
# For spark >= 4.0
|
|
from pyspark.errors.exceptions.base import IllegalArgumentException
|
|
except ModuleNotFoundError:
|
|
from pyspark.sql.utils import IllegalArgumentException
|
|
from pyspark.ml.base import Transformer
|
|
from pyspark.ml.functions import vector_to_array
|
|
from pyspark.ml.linalg import VectorUDT
|
|
from pyspark.ml.pipeline import PipelineModel
|
|
from pyspark.sql import DataFrame
|
|
from pyspark.sql import types as t
|
|
|
|
|
|
def cast_spark_df_with_vector_to_array(input_spark_df): # noqa: D417
|
|
"""
|
|
Finds columns of vector type in a spark dataframe and
|
|
casts them to array<double> type.
|
|
|
|
Args:
|
|
input_spark_df:
|
|
|
|
Returns:
|
|
A spark dataframe with vector columns transformed to array<double> type
|
|
"""
|
|
vector_type_columns = [
|
|
_field.name for _field in input_spark_df.schema if isinstance(_field.dataType, VectorUDT)
|
|
]
|
|
return reduce(
|
|
lambda df, vector_col: df.withColumn(vector_col, vector_to_array(vector_col)),
|
|
vector_type_columns,
|
|
input_spark_df,
|
|
)
|
|
|
|
|
|
def _do_pipeline_transform(df: DataFrame, transformer: Union[Transformer, PipelineModel]): # noqa: D417
|
|
"""
|
|
A util method that runs transform on a pipeline model/transformer
|
|
|
|
Args:
|
|
df: a spark dataframe
|
|
|
|
Returns:
|
|
output transformed dataframe using pipeline model/transformer
|
|
"""
|
|
return transformer.transform(df)
|
|
|
|
|
|
def _get_struct_type_by_cols(input_fields: set[str], df_schema: t.StructType) -> t.StructType:
|
|
"""
|
|
Args:
|
|
input_fields: A set of input columns to be
|
|
intersected with the input dataset's columns.
|
|
df_schema: A Spark dataframe schema to compare input_fields
|
|
|
|
Returns:
|
|
A StructType from the intersection of given columns and
|
|
the columns present in the training dataset
|
|
"""
|
|
if len(input_fields) > 0:
|
|
return t.StructType([_field for _field in df_schema.fields if _field.name in input_fields])
|
|
return []
|
|
|
|
|
|
def get_feature_cols(
|
|
df: DataFrame,
|
|
transformer: Union[Transformer, PipelineModel],
|
|
) -> set[str]:
|
|
"""
|
|
Finds feature columns from an input dataset. If a dataset
|
|
contains non-feature columns, those columns are not returned, but
|
|
if `input_fields` is set to include non-feature columns those
|
|
will be included in the return set of column names.
|
|
|
|
Args:
|
|
df: An input spark dataframe.
|
|
transformer: A pipeline/transformer to get the required feature columns
|
|
|
|
Returns:
|
|
A set of all the feature columns that are required
|
|
for the pipeline/transformer plus any initial columns passed in.
|
|
"""
|
|
feature_cols = set()
|
|
df_subset = df.limit(1).cache()
|
|
for column in df.columns:
|
|
try:
|
|
transformer.transform(df_subset.drop(column))
|
|
except IllegalArgumentException as iae:
|
|
if re.search("does not exist|no such struct field", str(iae), re.IGNORECASE):
|
|
feature_cols.add(column)
|
|
continue
|
|
raise
|
|
df_subset.unpersist()
|
|
return feature_cols
|