<commit_before>package shark.execution

import java.util.{HashMap => JHashMap, List => JList}

import org.apache.hadoop.io.BytesWritable

import org.apache.hadoop.hive.ql.exec.{ExprNodeEvaluator, JoinUtil}
import org.apache.hadoop.hive.ql.exec.HashTableSinkOperator.{HashTableSinkObjectCtx => MapJoinObjectCtx}
import org.apache.hadoop.hive.ql.exec.MapJoinMetaData
import org.apache.hadoop.hive.ql.exec.{MapJoinOperator => HiveMapJoinOperator}
import org.apache.hadoop.hive.ql.exec.persistence.{AbstractMapJoinKey, MapJoinDoubleKeys}
import org.apache.hadoop.hive.ql.exec.persistence.{MapJoinObjectKey, MapJoinSingleKey}
import org.apache.hadoop.hive.ql.exec.persistence.{MapJoinRowContainer, MapJoinObjectValue}
import org.apache.hadoop.hive.ql.plan.MapJoinDesc
import org.apache.hadoop.hive.ql.plan.{PartitionDesc, TableDesc}
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils.ObjectInspectorCopyOption
import org.apache.hadoop.hive.serde2.SerDe

import scala.collection.JavaConversions._
import scala.reflect.BeanProperty

import spark.broadcast.Broadcast
import spark.RDD
import shark.SharkEnv

object MapJoinOperator {
  type MapJoinHashTable = JHashMap[AbstractMapJoinKey, MapJoinObjectValue]
}

/**
 * A join operator optimized for joining a large table with a number of small
 * tables that fit in memory. The join can be performed as a map only job that
 * avoids an expensive shuffle process.
 *
 * Different from Hive, we don't spill the hash tables to disk. If the "small"
 * tables are too big to fit in memory, the normal join should be used anyway.
 */
class MapJoinOperator extends CommonJoinOperator[MapJoinDesc, HiveMapJoinOperator] {

  @BeanProperty var posBigTable: Int = _
  @BeanProperty var bigTableAlias: Int = _
  @BeanProperty var bigTableAliasByte: java.lang.Byte = _

  @transient var joinKeys: JHashMap[java.lang.Byte, JList[ExprNodeEvaluator]] = _
  @transient var joinKeysObjectInspectors: JHashMap[java.lang.Byte, JList[ObjectInspector]] = _

  @transient val metadataKeyTag = -1
  @transient var joinValues: JHashMap[java.lang.Byte, JList[ExprNodeEvaluator]] = _

  override def initializeOnMaster() {
    super.initializeOnMaster()
    posBigTable = conf.getPosBigTable()
    bigTableAlias = order(posBigTable).toInt
    bigTableAliasByte = bigTableAlias.toByte

    // Also call initialize on slave since we want the joinKeys and joinVals to
    // be initialized so we can use them in combineMultipleRdds(). This also puts
    // serialization info for keys in MapJoinMetaData.
    initializeOnSlave()
  }

  override def initializeOnSlave() {
    super.initializeOnSlave()

    joinKeys = new JHashMap[java.lang.Byte, JList[ExprNodeEvaluator]]
    JoinUtil.populateJoinKeyValue(
      joinKeys, conf.getKeys(), order, CommonJoinOperator.NOTSKIPBIGTABLE)

    // A bit confusing but getObjectInspectorsFromEvaluators also initializes
    // the evaluators.
    joinKeysObjectInspectors = JoinUtil.getObjectInspectorsFromEvaluators(
      joinKeys, objectInspectors.toArray, CommonJoinOperator.NOTSKIPBIGTABLE)

    // Put serialization metadata for keys in MapJoinMetaData.
    setKeyMetaData()
  }

  override def execute(): RDD[_] = {
    val inputRdds = executeParents()
    combineMultipleRdds(inputRdds)
  }

  override def combineMultipleRdds(rdds: Seq[(Int, RDD[_])]): RDD[_] = {
    logInfo("%d small tables to map join a large table (%d)".format(rdds.size - 1, posBigTable))

    val op1 = OperatorSerializationWrapper(this)
 
    initializeOnSlave()

    // Build hash tables for the small tables.
    val hashtables = rdds.zipWithIndex.filter(_._2 != bigTableAlias).map { case ((_, rdd), pos) =>

      logInfo("Creating hash table for input %d".format(pos))

      // First compute the keys and values of the small RDDs on slaves.
      // We need to do this before collecting the RDD because the RDD might
      // contain lazy structs that cannot be properly collected directly.
      val posByte = pos.toByte

      // Put serialization metadata for values in master's MapJoinMetaData.
      // Needed to deserialize values in collect().
      setValueMetaData(posByte)

      // Create a local reference for the serialized arrays, otherwise the
      // following mapParititons will fail because it tries to include the
      // outer closure, which references "this".
      val op = op1
      val rddForHash: RDD[(AbstractMapJoinKey, MapJoinObjectValue)] =
        rdd.mapPartitions { partition =>
          op.initializeOnSlave()
          // Put serialization metadata for values in slave's MapJoinMetaData.
          // Needed to serialize values in collect().
          op.setValueMetaData(posByte)
          op.computeJoinKeyValuesOnPartition(partition, posByte)
        }

      // Collect the RDD and build a hash table.
      val startCollect = System.currentTimeMillis()
      val wrappedRows: Array[(AbstractMapJoinKey, MapJoinObjectValue)] = rddForHash.collect()
      val collectTime = System.currentTimeMillis() - startCollect

      // Build the hash table.
      val startHash = System.currentTimeMillis()
      val hashTable = new MapJoinOperator.MapJoinHashTable
      wrappedRows.foreach { case (wrappedKey, wrappedValue) =>
        var mapEntry = hashTable.get(wrappedKey)
        if (mapEntry == null) {
          val container = new MapJoinRowContainer[Array[Object]]
          mapEntry = new MapJoinObjectValue(posByte, container)
          hashTable.put(wrappedKey, mapEntry)
        }
        mapEntry.getObj().add(wrappedValue.getObj.getList().get(0))
      }
      val hashTime = System.currentTimeMillis() - startHash
      logInfo("Input %d (%d rows) took %d ms to collect and %s ms to build hash table.".format(<commit_msg>Fixed a bug in map join that some rows are missing.<commit_after>package shark.execution

import java.util.{HashMap => JHashMap, List => JList}

import org.apache.hadoop.io.BytesWritable

import org.apache.hadoop.hive.ql.exec.{ExprNodeEvaluator, JoinUtil}
import org.apache.hadoop.hive.ql.exec.HashTableSinkOperator.{HashTableSinkObjectCtx => MapJoinObjectCtx}
import org.apache.hadoop.hive.ql.exec.MapJoinMetaData
import org.apache.hadoop.hive.ql.exec.{MapJoinOperator => HiveMapJoinOperator}
import org.apache.hadoop.hive.ql.exec.persistence.{AbstractMapJoinKey, MapJoinDoubleKeys}
import org.apache.hadoop.hive.ql.exec.persistence.{MapJoinObjectKey, MapJoinSingleKey}
import org.apache.hadoop.hive.ql.exec.persistence.{MapJoinRowContainer, MapJoinObjectValue}
import org.apache.hadoop.hive.ql.plan.MapJoinDesc
import org.apache.hadoop.hive.ql.plan.{PartitionDesc, TableDesc}
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorUtils.ObjectInspectorCopyOption
import org.apache.hadoop.hive.serde2.SerDe

import scala.collection.JavaConversions._
import scala.reflect.BeanProperty

import spark.broadcast.Broadcast
import spark.RDD
import shark.SharkEnv

object MapJoinOperator {
  type MapJoinHashTable = JHashMap[AbstractMapJoinKey, MapJoinObjectValue]
}

/**
 * A join operator optimized for joining a large table with a number of small
 * tables that fit in memory. The join can be performed as a map only job that
 * avoids an expensive shuffle process.
 *
 * Different from Hive, we don't spill the hash tables to disk. If the "small"
 * tables are too big to fit in memory, the normal join should be used anyway.
 */
class MapJoinOperator extends CommonJoinOperator[MapJoinDesc, HiveMapJoinOperator] {

  @BeanProperty var posBigTable: Int = _
  @BeanProperty var bigTableAlias: Int = _
  @BeanProperty var bigTableAliasByte: java.lang.Byte = _

  @transient var joinKeys: JHashMap[java.lang.Byte, JList[ExprNodeEvaluator]] = _
  @transient var joinKeysObjectInspectors: JHashMap[java.lang.Byte, JList[ObjectInspector]] = _

  @transient val metadataKeyTag = -1
  @transient var joinValues: JHashMap[java.lang.Byte, JList[ExprNodeEvaluator]] = _

  override def initializeOnMaster() {
    super.initializeOnMaster()
    posBigTable = conf.getPosBigTable()
    bigTableAlias = order(posBigTable).toInt
    bigTableAliasByte = bigTableAlias.toByte

    // Also call initialize on slave since we want the joinKeys and joinVals to
    // be initialized so we can use them in combineMultipleRdds(). This also puts
    // serialization info for keys in MapJoinMetaData.
    initializeOnSlave()
  }

  override def initializeOnSlave() {
    super.initializeOnSlave()

    joinKeys = new JHashMap[java.lang.Byte, JList[ExprNodeEvaluator]]
    JoinUtil.populateJoinKeyValue(
      joinKeys, conf.getKeys(), order, CommonJoinOperator.NOTSKIPBIGTABLE)

    // A bit confusing but getObjectInspectorsFromEvaluators also initializes
    // the evaluators.
    joinKeysObjectInspectors = JoinUtil.getObjectInspectorsFromEvaluators(
      joinKeys, objectInspectors.toArray, CommonJoinOperator.NOTSKIPBIGTABLE)

    // Put serialization metadata for keys in MapJoinMetaData.
    setKeyMetaData()
  }

  override def execute(): RDD[_] = {
    val inputRdds = executeParents()
    combineMultipleRdds(inputRdds)
  }

  override def combineMultipleRdds(rdds: Seq[(Int, RDD[_])]): RDD[_] = {
    logInfo("%d small tables to map join a large table (%d)".format(rdds.size - 1, posBigTable))

    val op1 = OperatorSerializationWrapper(this)
 
    initializeOnSlave()

    // Build hash tables for the small tables.
    val hashtables = rdds.zipWithIndex.filter(_._2 != bigTableAlias).map { case ((_, rdd), pos) =>

      logInfo("Creating hash table for input %d".format(pos))

      // First compute the keys and values of the small RDDs on slaves.
      // We need to do this before collecting the RDD because the RDD might
      // contain lazy structs that cannot be properly collected directly.
      val posByte = pos.toByte

      // Put serialization metadata for values in master's MapJoinMetaData.
      // Needed to deserialize values in collect().
      setValueMetaData(posByte)

      // Create a local reference for the serialized arrays, otherwise the
      // following mapParititons will fail because it tries to include the
      // outer closure, which references "this".
      val op = op1
      val rddForHash: RDD[(AbstractMapJoinKey, MapJoinObjectValue)] =
        rdd.mapPartitions { partition =>
          op.initializeOnSlave()
          // Put serialization metadata for values in slave's MapJoinMetaData.
          // Needed to serialize values in collect().
          op.setValueMetaData(posByte)
          op.computeJoinKeyValuesOnPartition(partition, posByte)
        }

      // Collect the RDD and build a hash table.
      val startCollect = System.currentTimeMillis()
      val wrappedRows: Array[(AbstractMapJoinKey, MapJoinObjectValue)] = rddForHash.collect()
      val collectTime = System.currentTimeMillis() - startCollect

      // Build the hash table.
      val startHash = System.currentTimeMillis()
      val hashTable = new MapJoinOperator.MapJoinHashTable
      wrappedRows.foreach { case (wrappedKey, wrappedValue) =>
        var mapEntry = hashTable.get(wrappedKey)
        if (mapEntry == null) {
          val container = new MapJoinRowContainer[Array[Object]]
          mapEntry = new MapJoinObjectValue(posByte, container)
          hashTable.put(wrappedKey, mapEntry)
          container.setList(wrappedValue.getObj.getList())
        } else {
          wrappedValue.getObj.getList().foreach(mapEntry.getObj().add)
        }
      }
      val hashTime = System.currentTimeMillis() - startHash
      logInfo("Input %d (%d rows) took %d ms to collect and %s ms to build hash table.".format(
{'fraction_non_alphanumeric': 0.08640639307416965, 'fraction_numerical': 0.0017481062182635478, 'mean_word_length': 3.8817553839902477, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 17, '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': '4148625', 'n_tokens_mistral': 3494, 'n_tokens_neox': 3475, 'n_words': 939}