def _transform(self, dataset):
    """Transforms the input DataFrame by applying the _run_model() mapPartitions function.

    Args:
      :dataset: A Spark DataFrame for TensorFlow inferencing.
    """
    spark = SparkSession.builder.getOrCreate()

    # set a deterministic order for input/output columns (lexicographic by key)
    input_cols = [col for col, tensor in sorted(self.getInputMapping().items())]      # input col => input tensor
    output_cols = [col for tensor, col in sorted(self.getOutputMapping().items())]    # output tensor => output col

    # run single-node inferencing on each executor
    logging.info("input_cols: {}".format(input_cols))
    logging.info("output_cols: {}".format(output_cols))

    # merge args + params
    logging.info("===== 1. inference args: {0}".format(self.args))
    logging.info("===== 2. inference params: {0}".format(self._paramMap))
    local_args = self.merge_args_params()
    logging.info("===== 3. inference args + params: {0}".format(local_args))

    tf_args = self.args.argv if self.args.argv else local_args
    rdd_out = dataset.select(input_cols).rdd.mapPartitions(lambda it: _run_model(it, local_args, tf_args))

    # convert to a DataFrame-friendly format
    rows_out = rdd_out.map(lambda x: Row(*x))
    return spark.createDataFrame(rows_out, output_cols)
def _transform ( self , dataset ) : spark = SparkSession . builder . getOrCreate ( ) input_cols = [ col for col , tensor in sorted ( self . getInputMapping ( ) . items ( ) ) ] output_cols = [ col for tensor , col in sorted ( self . getOutputMapping ( ) . items ( ) ) ] logging . info ( "input_cols: {}" . format ( input_cols ) ) logging . info ( "output_cols: {}" . format ( output_cols ) ) logging . info ( "===== 1. inference args: {0}" . format ( self . args ) ) logging . info ( "===== 2. inference params: {0}" . format ( self . _paramMap ) ) local_args = self . merge_args_params ( ) logging . info ( "===== 3. inference args + params: {0}" . format ( local_args ) ) tf_args = self . args . argv if self . args . argv else local_args rdd_out = dataset . select ( input_cols ) . rdd . mapPartitions ( lambda it : _run_model ( it , local_args , tf_args ) ) rows_out = rdd_out . map ( lambda x : Row ( * x ) ) return spark . createDataFrame ( rows_out , output_cols )
Transforms the input DataFrame by applying the _run_model() mapPartitions function.

    Args:
      :dataset: A Spark DataFrame for TensorFlow inferencing.