<commit_before>/*
 * Copyright (c) 2016 SnappyData, Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you
 * may not use this file except in compliance with the License. You
 * may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License. See accompanying
 * LICENSE file.
 */
package org.apache.spark.sql.execution.joins

import java.io.{ObjectOutputStream, IOException}

import org.apache.spark.annotation.DeveloperApi
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.catalyst.plans.physical.{Distribution, Partitioning, UnspecifiedDistribution}
import org.apache.spark.sql.execution.metric.SQLMetrics
import org.apache.spark.sql.execution.{BinaryNode, SparkPlan}
import org.apache.spark.util.Utils
import org.apache.spark.{OneToOneDependency, Partition, SparkContext, TaskContext}


/**
 * :: DeveloperApi ::
 * Performs an local hash join of two child relations.  If a relation (out of a datasource) is already replicated
 * accross all nodes then rather than doing a Broadcast join which can be expensive, this join just
 * scans through the single partition of the replicated relation while streaming through the other relation
 */
@DeveloperApi
case class LocalJoin(leftKeys: Seq[Expression],
    rightKeys: Seq[Expression],
    buildSide: BuildSide,
    left: SparkPlan,
    right: SparkPlan)
    extends BinaryNode with HashJoin {


  override private[sql] lazy val metrics = Map(
    "numLeftRows" -> SQLMetrics.createLongMetric(sparkContext, "number of left rows"),
    "numRightRows" -> SQLMetrics.createLongMetric(sparkContext, "number of right rows"),
    "numOutputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of output rows"))

  override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning


  override def requiredChildDistribution: Seq[Distribution] =
    UnspecifiedDistribution :: UnspecifiedDistribution :: Nil

  /**
   * Overridden by concrete implementations of SparkPlan.
   * Produces the result of the query as an RDD[InternalRow]
   */
  override protected def doExecute(): RDD[InternalRow] = {
    val (numBuildRows, numStreamedRows) = buildSide match {
      case BuildLeft => (longMetric("numLeftRows"), longMetric("numRightRows"))
      case BuildRight => (longMetric("numRightRows"), longMetric("numLeftRows"))
    }
    val numOutputRows = longMetric("numOutputRows")

    narrowPartitions(buildPlan.execute(), streamedPlan.execute(),
      preservesPartitioning = true) {
      (buildIter, streamIter) => {
        val hashed = HashedRelation(buildIter, numBuildRows, buildSideKeyGenerator)
        hashJoin(streamIter, numStreamedRows, hashed, numOutputRows)
      }
    }
  }

  def narrowPartitions(buildRDD: RDD[InternalRow], streamRDD: RDD[InternalRow], preservesPartitioning: Boolean)
      (f: (Iterator[InternalRow], Iterator[InternalRow]) => Iterator[InternalRow]): NarrowPartitionsRDD = {
    val sc = buildRDD.sparkContext
    new NarrowPartitionsRDD(sc, sc.clean(f),
      buildRDD, streamRDD, preservesPartitioning)
  }
}


private[spark] class NarrowPartitionsRDD(
    @transient sc: SparkContext,
    var f: (Iterator[InternalRow], Iterator[InternalRow]) => Iterator[InternalRow],
    var buildRDD: RDD[InternalRow],
    var streamRDD: RDD[InternalRow],
    preservesPartitioning: Boolean = false)
    extends RDD[InternalRow](sc, Seq(new OneToOneDependency(streamRDD))) {

  override def compute(s: Partition, context: TaskContext): Iterator[InternalRow] = {
    val partitions = s.asInstanceOf[NarrowPartitionsPartition]
    f(buildRDD.iterator(partitions.buildPartition, context),
      streamRDD.iterator(partitions.streamPartition, context))
  }

  override def getPartitions: Array[Partition] = {
    val numParts = streamRDD.partitions.length
    val part = buildRDD.partitions.head
    Array.tabulate[Partition](numParts) { i =>
      val streamLocs = streamRDD.preferredLocations(streamRDD.partitions(i))
      val buildLocs = buildRDD.preferredLocations(part)
      val exactMatchLocations = streamLocs.intersect(buildLocs)
      val locs = if (exactMatchLocations.nonEmpty) exactMatchLocations else
                         (streamLocs ++ buildLocs).distinct

      new NarrowPartitionsPartition(part.index, buildRDD, i, streamRDD , locs)
    }
  }

  override def getPreferredLocations(s: Partition): Seq[String] = {
    s.asInstanceOf[NarrowPartitionsPartition].preferredLocations
  }

  override def clearDependencies() {
    super.clearDependencies()
    buildRDD = null
    streamRDD = null
    f = null
  }

}

private[spark] class NarrowPartitionsPartition(
    buildIdx : Int,
    @transient var buildRDD: RDD[InternalRow],
    streamIdx: Int,
    @transient var streamRDD: RDD[InternalRow],
    @transient val preferredLocations: Seq[String])
    extends Partition {
  override val index: Int = streamIdx
  var buildPartition = buildRDD.partitions(buildIdx)
  var streamPartition = streamRDD.partitions(streamIdx)

  @throws[IOException]
  private def writeObject(oos: ObjectOutputStream): Unit = Utils.tryOrIOException {
    // Update the reference to parent split at the time of task serialization
    buildPartition = buildRDD.partitions(buildIdx)
    streamPartition = streamRDD.partitions(streamIdx)

    oos.defaultWriteObject()
  }
}
<commit_msg>Snap 634 (#206)<commit_after>/*
 * Copyright (c) 2016 SnappyData, Inc. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you
 * may not use this file except in compliance with the License. You
 * may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * permissions and limitations under the License. See accompanying
 * LICENSE file.
 */
package org.apache.spark.sql.execution.joins

import java.io.{ObjectOutputStream, IOException}


import scala.collection.concurrent.TrieMap
import scala.collection.mutable.HashMap

import com.google.common.cache.{CacheLoader, CacheBuilder}

import org.apache.spark.annotation.DeveloperApi
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.Expression
import org.apache.spark.sql.catalyst.plans.physical.{Distribution, Partitioning, UnspecifiedDistribution}
import org.apache.spark.sql.execution.metric.SQLMetrics
import org.apache.spark.sql.execution.{BinaryNode, SparkPlan}
import org.apache.spark.storage.{StorageLevel, RDDBlockId, BroadcastBlockId}
import org.apache.spark.util.Utils
import org.apache.spark.{SparkEnv, OneToOneDependency, Partition, SparkContext, TaskContext}


/**
 * :: DeveloperApi ::
 * Performs an local hash join of two child relations.  If a relation (out of a datasource) is already replicated
 * accross all nodes then rather than doing a Broadcast join which can be expensive, this join just
 * scans through the single partition of the replicated relation while streaming through the other relation
 */
@DeveloperApi
case class LocalJoin(leftKeys: Seq[Expression],
    rightKeys: Seq[Expression],
    buildSide: BuildSide,
    left: SparkPlan,
    right: SparkPlan)
    extends BinaryNode with HashJoin {


  override private[sql] lazy val metrics = Map(
    "numLeftRows" -> SQLMetrics.createLongMetric(sparkContext, "number of left rows"),
    "numRightRows" -> SQLMetrics.createLongMetric(sparkContext, "number of right rows"),
    "numOutputRows" -> SQLMetrics.createLongMetric(sparkContext, "number of output rows"))

  override def outputPartitioning: Partitioning = streamedPlan.outputPartitioning


  override def requiredChildDistribution: Seq[Distribution] =
    UnspecifiedDistribution :: UnspecifiedDistribution :: Nil

  /**
   * Overridden by concrete implementations of SparkPlan.
   * Produces the result of the query as an RDD[InternalRow]
   */
  override protected def doExecute(): RDD[InternalRow] = {
    val (numBuildRows, numStreamedRows) = buildSide match {
      case BuildLeft => (longMetric("numLeftRows"), longMetric("numRightRows"))
      case BuildRight => (longMetric("numRightRows"), longMetric("numLeftRows"))
    }
    val numOutputRows = longMetric("numOutputRows")

    def hashedRelationIter(buildIter: Iterator[InternalRow]): HashedRelation = {
      HashedRelation(buildIter, numBuildRows, buildSideKeyGenerator)
    }

    val buildRDD = buildPlan.execute()
    val streamRDD = streamedPlan.execute()

    val sc = buildRDD.sparkContext
    val hashedRDD = new HashRelationRDD(sc, buildRDD,
      streamRDD.partitions.length, sc.clean(hashedRelationIter))

    narrowPartitions(hashedRDD, streamRDD, true) {
      (hashedIter, streamIter) => {
        val hashed = hashedIter.next()
        hashJoin(streamIter, numStreamedRows, hashed, numOutputRows)
      }
    }
  }

  def narrowPartitions(hashedRDD: RDD[HashedRelation], streamRDD: RDD[InternalRow],
      preservesPartitioning: Boolean)
      (f: (Iterator[HashedRelation], Iterator[InternalRow])
          => Iterator[InternalRow]): NarrowPartitionsRDD = {
    val sc = hashedRDD.sparkContext
    new NarrowPartitionsRDD(sc, sc.clean(f),
      hashedRDD, streamRDD, preservesPartitioning)
  }
}

// Helper object jut to take a global sync across tasks.
// Furture helper methods can come here
private[spark] object HashRelationRDD

private[spark] class HashRelationRDD(
    sc: SparkContext,
    var buildRDD: RDD[InternalRow],
    val maxPartitions: Int,
    var f: (Iterator[InternalRow]) => HashedRelation
    ) extends RDD[HashedRelation](sc, Seq(new OneToOneDependency(buildRDD))) {

  override def compute(s: Partition, context: TaskContext): Iterator[HashedRelation] = {

    val blockId = RDDBlockId(this.id, s.index)

    val rel1 = SparkEnv.get.blockManager.getSingle(blockId) match {
      case Some(x) => x.asInstanceOf[HashedRelation]
      case None => {
        HashRelationRDD.synchronized {
          SparkEnv.get.blockManager.getSingle(blockId) match {
            case Some(x) => x.asInstanceOf[HashedRelation]
            case None => {
              val hashedRelation = f(buildRDD.iterator(s, context))
              SparkEnv.get.blockManager.putSingle(
                blockId, hashedRelation, StorageLevel.MEMORY_AND_DISK_SER, tellMaster = false)

              hashedRelation
            }
          }

        }
      }
    }

    Seq(rel1).iterator
  }

  override def getPartitions: Array[Partition] = {
    buildRDD.partitions
  }


  override def getPreferredLocations(s: Partition): Seq[String] = {
    buildRDD.preferredLocations(s)
  }
}

private[spark] class NarrowPartitionsRDD(
    @transient sc: SparkContext,
    var f: (Iterator[HashedRelation], Iterator[InternalRow]) => Iterator[InternalRow],
    var hashedRDD: RDD[HashedRelation],
    var streamRDD: RDD[InternalRow],
    preservesPartitioning: Boolean = false)
    extends RDD[InternalRow](sc, Seq(new OneToOneDependency(streamRDD))) {

  override def compute(s: Partition, context: TaskContext): Iterator[InternalRow] = {
    val partitions = s.asInstanceOf[NarrowPartitionsPartition]
    f(hashedRDD.iterator(partitions.buildPartition, context),
      streamRDD.iterator(partitions.streamPartition, context))
  }

  override def getPartitions: Array[Partition] = {
    val numParts = streamRDD.partitions.length
    val part = hashedRDD.partitions.head
    Array.tabulate[Partition](numParts) { i =>
      val streamLocs = streamRDD.preferredLocations(streamRDD.partitions(i))
      val buildLocs = hashedRDD.preferredLocations(part)
      val exactMatchLocations = streamLocs.intersect(buildLocs)
      val locs = if (exactMatchLocations.nonEmpty) exactMatchLocations else
                         (streamLocs ++ buildLocs).distinct

      new NarrowPartitionsPartition(part.index, hashedRDD, i, streamRDD , locs)
    }
  }

  override def getPreferredLocations(s: Partition): Seq[String] = {
    s.asInstanceOf[NarrowPartitionsPartition].preferredLocations
  }

  override def clearDependencies() {
    super.clearDependencies()
    hashedRDD = null
    streamRDD = null
    f = null
  }

}

private[spark] class NarrowPartitionsPartition(
    buildIdx: Int,
    @transient var hashedRDD: RDD[HashedRelation],
    streamIdx: Int,
    @transient var streamRDD: RDD[InternalRow],
    @transient val preferredLocations: Seq[String])
    extends Partition {
  override val index: Int = streamIdx
  var buildPartition = hashedRDD.partitions(buildIdx)
  var streamPartition = streamRDD.partitions(streamIdx)

  @throws[IOException]
  private def writeObject(oos: ObjectOutputStream): Unit = Utils.tryOrIOException {
    // Update the reference to parent split at the time of task serialization
    buildPartition = hashedRDD.partitions(buildIdx)
    streamPartition = streamRDD.partitions(streamIdx)

    oos.defaultWriteObject()
  }
}

{'fraction_non_alphanumeric': 0.0917784861739195, 'fraction_numerical': 0.0017792275187189562, 'mean_word_length': 4.275713727023856, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 26, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5327058', 'n_tokens_mistral': 4051, 'n_tokens_neox': 3875, 'n_words': 1007}