Source code for synapse.ml.cyber.feature.indexers

__author__ = "rolevin"

from typing import List

from synapse.ml.cyber.utils.spark_utils import (
    DataFrameUtils,
    ExplainBuilder,
    HasSetInputCol,
    HasSetOutputCol,
)

from pyspark.ml import Estimator, Transformer
from pyspark.ml.param.shared import HasInputCol, HasOutputCol, Param, Params
from pyspark.sql import DataFrame, functions as f


class IdIndexerModel(Transformer, HasSetInputCol, HasSetOutputCol):
    partitionKey = Param(
        Params._dummy(),
        "partitionKey",
        "The name of the column to partition by, i.e., make sure the indexing takes the partition into account. "
        "This is exemplified in reset_per_partition.",
    )

    def __init__(
        self,
        input_col: str,
        partition_key: str,
        output_col: str,
        vocab_df: DataFrame,
    ):
        super().__init__()
        ExplainBuilder.build(
            self,
            inputCol=input_col,
            partitionKey=partition_key,
            outputCol=output_col,
        )
        self._vocab_df = vocab_df

    def undo_transform(self, df: DataFrame) -> DataFrame:
        ucols = [self.partition_key, self.output_col]
        vocab_df = self._vocab_df

        return df.join(vocab_df, on=ucols, how="left_outer")

    def _transform(self, df):
        ucols = [self.partition_key, self.input_col]

        input_col = self.input_col
        output_col = self.output_col
        vocab_df = self._vocab_df

        return (
            df.join(vocab_df, on=ucols, how="left_outer")
            .withColumn(
                output_col,
                f.when(f.col(output_col).isNotNull(), f.col(output_col)).otherwise(
                    f.lit(0),
                ),
            )
            .drop(input_col)
        )


class IdIndexer(Estimator, HasSetInputCol, HasSetOutputCol):
    partitionKey = Param(
        Params._dummy(),
        "partitionKey",
        "The name of the column to partition by, i.e., make sure the indexing takes the partition into account. "
        "This is exemplified in reset_per_partition.",
    )

    resetPerPartition = Param(
        Params._dummy(),
        "resetPerPartition",
        "When set to True then indexing is consecutive from [1..n] for each value of the partition column. "
        "When set to False then indexing is consecutive for all partition and column values.",
    )

    def __init__(
        self,
        input_col: str,
        partition_key: str,
        output_col: str,
        reset_per_partition: bool,
    ):
        super().__init__()

        ExplainBuilder.build(
            self,
            inputCol=input_col,
            partitionKey=partition_key,
            outputCol=output_col,
            resetPerPartition=reset_per_partition,
        )

    def _make_vocab_df(self, df):
        ucols = [self.getPartitionKey(), self.getInputCol()]
        the_df = df.select(ucols).distinct().orderBy(ucols)

        return (
            DataFrameUtils.zip_with_index(
                df=the_df,
                start_index=1,
                col_name=self.getOutputCol(),
                partition_col=self.getPartitionKey(),
                order_by_col=self.getInputCol(),
            )
            if self.getResetPerPartition()
            else DataFrameUtils.zip_with_index(
                df=the_df,
                start_index=1,
                col_name=self.getOutputCol(),
            )
        )

    def _fit(self, df: DataFrame) -> IdIndexerModel:
        return IdIndexerModel(
            self.input_col,
            self.partition_key,
            self.output_col,
            self._make_vocab_df(df).cache(),
        )


class MultiIndexerModel(Transformer):
    def __init__(self, models: List[IdIndexerModel]):
        super().__init__()
        self.models = models

    def get_model_by_input_col(self, input_col):
        for m in self.models:
            if m.input_col == input_col:
                return m

        return None

    def get_model_by_output_col(self, output_col):
        for m in self.models:
            if m.output_col == output_col:
                return m

        return None

    def undo_transform(self, df: DataFrame) -> DataFrame:
        curr_df = df.cache()

        for model in self.models:
            curr_df = model.undo_transform(curr_df).cache()

        return curr_df

    def _transform(self, df: DataFrame) -> DataFrame:
        curr_df = df.cache()

        for model in self.models:
            curr_df = model.transform(curr_df).cache()

        return curr_df


class MultiIndexer(Estimator):
    def __init__(self, indexers: List[IdIndexer]):
        super().__init__()
        self.indexers = indexers

    def _fit(self, df: DataFrame) -> MultiIndexerModel:
        return MultiIndexerModel([i.fit(df) for i in self.indexers])