star

venv/lib/python3.9/site-packages/mlflow/pyspark/ml/_autolog.py

@@ -0,0 +1,97 @@
+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

venv/lib/python3.9/site-packages/mlflow/pyspark/ml/log_model_allowlist.txt

@@ -0,0 +1,62 @@
+# classification
+pyspark.ml.classification.LinearSVCModel
+pyspark.ml.classification.DecisionTreeClassificationModel
+pyspark.ml.classification.GBTClassificationModel
+pyspark.ml.classification.LogisticRegressionModel
+pyspark.ml.classification.RandomForestClassificationModel
+pyspark.ml.classification.NaiveBayesModel
+
+# clustering
+pyspark.ml.clustering.BisectingKMeansModel
+pyspark.ml.clustering.KMeansModel
+pyspark.ml.clustering.GaussianMixtureModel
+
+# Regression
+pyspark.ml.regression.AFTSurvivalRegressionModel
+pyspark.ml.regression.DecisionTreeRegressionModel
+pyspark.ml.regression.GBTRegressionModel
+pyspark.ml.regression.GeneralizedLinearRegressionModel
+pyspark.ml.regression.LinearRegressionModel
+pyspark.ml.regression.RandomForestRegressionModel
+
+# Featurizer model
+pyspark.ml.feature.BucketedRandomProjectionLSHModel
+pyspark.ml.feature.ChiSqSelectorModel
+pyspark.ml.feature.CountVectorizerModel
+pyspark.ml.feature.IDFModel
+pyspark.ml.feature.ImputerModel
+pyspark.ml.feature.MaxAbsScalerModel
+pyspark.ml.feature.MinHashLSHModel
+pyspark.ml.feature.MinMaxScalerModel
+pyspark.ml.feature.OneHotEncoderModel
+pyspark.ml.feature.RobustScalerModel
+pyspark.ml.feature.RFormulaModel
+pyspark.ml.feature.StandardScalerModel
+pyspark.ml.feature.StringIndexerModel
+pyspark.ml.feature.VarianceThresholdSelectorModel
+pyspark.ml.feature.VectorIndexerModel
+pyspark.ml.feature.UnivariateFeatureSelectorModel
+
+# composite model
+pyspark.ml.classification.OneVsRestModel
+
+# pipeline model
+pyspark.ml.pipeline.PipelineModel
+
+# Hyper-parameter tuning
+pyspark.ml.tuning.CrossValidatorModel
+pyspark.ml.tuning.TrainValidationSplitModel
+
+# SynapeML models
+synapse.ml.cognitive.*
+synapse.ml.exploratory.*
+synapse.ml.featurize.*
+synapse.ml.geospatial.*
+synapse.ml.image.*
+synapse.ml.io.*
+synapse.ml.isolationforest.*
+synapse.ml.lightgbm.*
+synapse.ml.nn.*
+synapse.ml.opencv.*
+synapse.ml.stages.*
+synapse.ml.vw.*