<commit_before>package workflow.graph

import org.apache.spark.rdd.RDD

import scala.reflect.ClassTag

/**
 * A Pipeline takes data as input (single item or an RDD), and outputs some transformation
 * of that data. Internally, a Pipeline contains a [[GraphExecutor]], a specified source, and a specified sink.
 * When a pipeline is applied to data it produces a [[GraphExecution]], in the form of either a [[PipelineDatasetOut]]
 * or a [[PipelineDatumOut]]. These are lazy wrappers around the scheduled execution under the hood,
 * and when their values are accessed the underlying [[Graph]] will be executed.
 *
 * Warning: Not thread-safe!
 *
 * @tparam A type of the data this Pipeline expects as input
 * @tparam B type of the data this Pipeline outputs
 */
trait Pipeline[A, B] {
  private[graph] val source: SourceId
  private[graph] val sink: SinkId
  private[graph] def executor: GraphExecutor

  /**
   * Lazily apply the pipeline to a single datum.
   * @return A lazy wrapper around the result of passing the datum through the pipeline.
   */
  final def apply(datum: A): PipelineDatumOut[B] = {
    new PipelineDatumOut[B](executor, sink, Some(source, datum))
  }

  /**
   * Lazily apply the pipeline to a dataset.
   * @return A lazy wrapper around the result of passing the dataset through the pipeline.
   */
  final def apply(data: RDD[A]): PipelineDatasetOut[B] = {
    new PipelineDatasetOut[B](executor, sink, Some(source, data))
  }

  /**
   * Lazily apply the pipeline to the lazy output of a different pipeline given an initial dataset.
   * If the previous pipeline has already been fit, it will not need to be fit again.
   * @return A lazy wrapper around the result of passing lazy output from a different pipeline through this pipeline.
   */
  final def apply(data: PipelineDatasetOut[A]): PipelineDatasetOut[B] = {
    val (newGraph, sourceMapping, nodeMapping, sinkMapping) =
      data.getGraph.connectGraph(executor.getGraph, Map(source -> data.getSink))
    val idMapping: Map[GraphId, GraphId] = sourceMapping ++ nodeMapping ++ sinkMapping
    val newState = data.getState - data.getSink ++ executor.getState.map(x => (idMapping(x._1), x._2))

    new PipelineDatasetOut[B](new GraphExecutor(newGraph, newState), sinkMapping(sink), None)
  }

  /**
   * Lazily apply the pipeline to the lazy output of a different pipeline given an initial datum.
   * If the previous pipeline has already been fit, it will not need to be fit again.
   * @return A lazy wrapper around the result of passing lazy output from a different pipeline through this pipeline.
   */
  final def apply(datum: PipelineDatumOut[A]): PipelineDatumOut[B] = {
    val (newGraph, sourceMapping, nodeMapping, sinkMapping) =
      datum.getGraph.connectGraph(executor.getGraph, Map(source -> datum.getSink))
    val idMapping: Map[GraphId, GraphId] = sourceMapping ++ nodeMapping ++ sinkMapping
    val newState = datum.getState - datum.getSink ++ executor.getState.map(x => (idMapping(x._1), x._2))

    new PipelineDatumOut[B](new GraphExecutor(newGraph, newState), sinkMapping(sink), None)
  }

  /**
   * Chains a pipeline onto the end of this one, producing a new pipeline.
   * If either this pipeline or the following has already been fit, it will not need to be fit again.
   *
   * @param next the pipeline to chain
   */
  final def andThen[C](next: Pipeline[B, C]): Pipeline[A, C] = {
    val (newGraph, sourceMapping, nodeMapping, sinkMapping) =
      executor.getGraph.connectGraph(next.executor.getGraph, Map(next.source -> sink))
    val idMapping: Map[GraphId, GraphId] = sourceMapping ++ nodeMapping ++ sinkMapping
    val newState = executor.getState - sink ++ next.executor.getState.map(x => (idMapping(x._1), x._2))
    new ConcretePipeline(new GraphExecutor(newGraph, newState), source, sinkMapping(next.sink))
  }

  final def andThen[C](est: Estimator[B, C], data: RDD[A]): Pipeline[A, C] = {
    this andThen est.fit(apply(data))
  }

  final def andThen[C](est: Estimator[B, C], data: PipelineDatasetOut[A]): Pipeline[A, C] = {
    this andThen est.fit(apply(data))
  }

  final def andThen[C, L](
    est: LabelEstimator[B, C, L],
    data: RDD[A],
    labels: RDD[L]
  ): Pipeline[A, C] = {
    this andThen est.fit(apply(data), labels)
  }

  final def andThen[C, L](
    est: LabelEstimator[B, C, L],
    data: PipelineDatasetOut[A],
    labels: RDD[L]
  ): Pipeline[A, C] = {
    this andThen est.fit(apply(data), labels)
  }

  final def andThen[C, L](
    est: LabelEstimator[B, C, L],
    data: RDD[A],
    labels: PipelineDatasetOut[L]
  ): Pipeline[A, C] = {
    this andThen est.fit(apply(data), labels)
  }
<commit_msg>more documentation<commit_after>package workflow.graph

import org.apache.spark.rdd.RDD

import scala.reflect.ClassTag

/**
 * A Pipeline takes data as input (single item or an RDD), and outputs some transformation
 * of that data. Internally, a Pipeline contains a [[GraphExecutor]], a specified source, and a specified sink.
 * When a pipeline is applied to data it produces a [[GraphExecution]], in the form of either a [[PipelineDatasetOut]]
 * or a [[PipelineDatumOut]]. These are lazy wrappers around the scheduled execution under the hood,
 * and when their values are accessed the underlying [[Graph]] will be executed.
 *
 * Warning: Not thread-safe!
 *
 * @tparam A type of the data this Pipeline expects as input
 * @tparam B type of the data this Pipeline outputs
 */
trait Pipeline[A, B] {
  private[graph] val source: SourceId
  private[graph] val sink: SinkId
  private[graph] def executor: GraphExecutor

  /**
   * Lazily apply the pipeline to a single datum.
   * @return A lazy wrapper around the result of passing the datum through the pipeline.
   */
  final def apply(datum: A): PipelineDatumOut[B] = {
    new PipelineDatumOut[B](executor, sink, Some(source, datum))
  }

  /**
   * Lazily apply the pipeline to a dataset.
   * @return A lazy wrapper around the result of passing the dataset through the pipeline.
   */
  final def apply(data: RDD[A]): PipelineDatasetOut[B] = {
    new PipelineDatasetOut[B](executor, sink, Some(source, data))
  }

  /**
   * Lazily apply the pipeline to the lazy output of a different pipeline given an initial dataset.
   * If the previous pipeline has already been fit, it will not need to be fit again.
   * @return A lazy wrapper around the result of passing lazy output from a different pipeline through this pipeline.
   */
  final def apply(data: PipelineDatasetOut[A]): PipelineDatasetOut[B] = {
    val (newGraph, sourceMapping, nodeMapping, sinkMapping) =
      data.getGraph.connectGraph(executor.getGraph, Map(source -> data.getSink))
    val idMapping: Map[GraphId, GraphId] = sourceMapping ++ nodeMapping ++ sinkMapping
    val newState = data.getState - data.getSink ++ executor.getState.map(x => (idMapping(x._1), x._2))

    new PipelineDatasetOut[B](new GraphExecutor(newGraph, newState), sinkMapping(sink), None)
  }

  /**
   * Lazily apply the pipeline to the lazy output of a different pipeline given an initial datum.
   * If the previous pipeline has already been fit, it will not need to be fit again.
   * @return A lazy wrapper around the result of passing lazy output from a different pipeline through this pipeline.
   */
  final def apply(datum: PipelineDatumOut[A]): PipelineDatumOut[B] = {
    val (newGraph, sourceMapping, nodeMapping, sinkMapping) =
      datum.getGraph.connectGraph(executor.getGraph, Map(source -> datum.getSink))
    val idMapping: Map[GraphId, GraphId] = sourceMapping ++ nodeMapping ++ sinkMapping
    val newState = datum.getState - datum.getSink ++ executor.getState.map(x => (idMapping(x._1), x._2))

    new PipelineDatumOut[B](new GraphExecutor(newGraph, newState), sinkMapping(sink), None)
  }

  /**
   * Chains a pipeline onto the end of this one, producing a new pipeline.
   * If either this pipeline or the following has already been executed, it will not need to be fit again.
   *
   * @param next the pipeline to chain
   */
  final def andThen[C](next: Pipeline[B, C]): Pipeline[A, C] = {
    val (newGraph, sourceMapping, nodeMapping, sinkMapping) =
      executor.getGraph.connectGraph(next.executor.getGraph, Map(next.source -> sink))
    val idMapping: Map[GraphId, GraphId] = sourceMapping ++ nodeMapping ++ sinkMapping
    val newState = executor.getState - sink ++ next.executor.getState.map(x => (idMapping(x._1), x._2))
    new ConcretePipeline(new GraphExecutor(newGraph, newState), source, sinkMapping(next.sink))
  }

  /**
   * Chains an estimator onto the end of this pipeline, producing a new pipeline.
   * If this pipeline has already been executed, it will not need to be fit again.
   *
   * @param est The estimator to chain onto the end of this pipeline
   * @param data The training data to use
   *             (the estimator will be fit on the result of passing this data through the current pipeline)
   */
  final def andThen[C](est: Estimator[B, C], data: RDD[A]): Pipeline[A, C] = {
    this andThen est.fit(apply(data))
  }

  /**
   * Chains an estimator onto the end of this pipeline, producing a new pipeline.
   * If this pipeline has already been executed, it will not need to be fit again.
   *
   * @param est The estimator to chain onto the end of this pipeline
   * @param data The training data to use
   *             (the estimator will be fit on the result of passing this data through the current pipeline)
   */
  final def andThen[C](est: Estimator[B, C], data: PipelineDatasetOut[A]): Pipeline[A, C] = {
    this andThen est.fit(apply(data))
  }

  /**
   * Chains a label estimator onto the end of this pipeline, producing a new pipeline.
   * If this pipeline has already been executed, it will not need to be fit again.
   *
   * @param est The estimator to chain onto the end of this pipeline
   * @param data The training data to use
   *             (the estimator will be fit on the result of passing this data through the current pipeline)
   * @param labels The labels to use when fitting the LabelEstimator. Must be zippable with the training data.
   */
  final def andThen[C, L](
    est: LabelEstimator[B, C, L],
    data: RDD[A],
    labels: RDD[L]
  ): Pipeline[A, C] = {
    this andThen est.fit(apply(data), labels)
  }

  /**
   * Chains a label estimator onto the end of this pipeline, producing a new pipeline.
   * If this pipeline has already been executed, it will not need to be fit again.
   *
   * @param est The estimator to chain onto the end of this pipeline
   * @param data The training data to use
   *             (the estimator will be fit on the result of passing this data through the current pipeline)
   * @param labels The labels to use when fitting the LabelEstimator. Must be zippable with the training data.
   */
  final def andThen[C, L](
    est: LabelEstimator[B, C, L],
    data: PipelineDatasetOut[A],
    labels: RDD[L]
  ): Pipeline[A, C] = {
    this andThen est.fit(apply(data), labels)
  }

  /**
   * Chains a label estimator onto the end of this pipeline, producing a new pipeline.
   * If this pipeline has already been executed, it will not need to be fit again.
   *
   * @param est The estimator to chain onto the end of this pipeline
   * @param data The training data to use
   *             (the estimator will be fit on the result of passing this data through the current pipeline)
   * @param labels The labels to use when fitting the LabelEstimator. Must be zippable with the training data.
   */
  final def andThen[C, L](
    est: LabelEstimator[B, C, L],
    data: RDD[A],
    labels: PipelineDatasetOut[L]
  ): Pipeline[A, C] = {
    this andThen est.fit(apply(data), labels)
  }

  /**
   * Chains a label estimator onto the end of this pipeline, producing a new pipeline.
   * If this pipeline has already been executed, it will not need to be fit again.
   *
   * @param est The estimator to chain onto the end of this pipeline
   * @param data The training data to use
   *             (the estimator will be fit on the result of passing this data through the current pipeline)
   * @param labels The labels to use when fitting the LabelEstimator. Must be zippable with the training data.
   */
{'fraction_non_alphanumeric': 0.09573770491803278, 'fraction_numerical': 0.0009836065573770492, 'mean_word_length': 3.590293453724605, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4895198', 'n_tokens_mistral': 3642, 'n_tokens_neox': 3598, 'n_words': 1449}