/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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.
 */

package org.apache.spark.sql.execution.joins

import java.io._

import com.esotericsoftware.kryo.{Kryo, KryoSerializable}
import com.esotericsoftware.kryo.io.{Input, Output}

import org.apache.spark.{SparkConf, SparkEnv, SparkException, SparkUnsupportedOperationException}
import org.apache.spark.internal.config.{BUFFER_PAGESIZE, MEMORY_OFFHEAP_ENABLED}
import org.apache.spark.memory._
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
import org.apache.spark.sql.catalyst.plans.physical.BroadcastMode
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.types.LongType
import org.apache.spark.unsafe.Platform
import org.apache.spark.unsafe.map.BytesToBytesMap
import org.apache.spark.util.{KnownSizeEstimation, Utils}

/**
 * Interface for a hashed relation by some key. Use [[HashedRelation.apply]] to create a concrete
 * object.
 */
private[execution] sealed trait HashedRelation extends KnownSizeEstimation {
  /**
   * Returns matched rows.
   *
   * Returns null if there is no matched rows.
   */
  def get(key: InternalRow): Iterator[InternalRow]

  /**
   * Returns matched rows for a key that has only one column with LongType.
   *
   * Returns null if there is no matched rows.
   */
  def get(key: Long): Iterator[InternalRow] = {
    throw SparkUnsupportedOperationException()
  }

  /**
   * Returns the matched single row.
   */
  def getValue(key: InternalRow): InternalRow

  /**
   * Returns the matched single row with key that have only one column of LongType.
   */
  def getValue(key: Long): InternalRow = {
    throw SparkUnsupportedOperationException()
  }

  /**
   * Returns an iterator for key index and matched rows.
   *
   * Returns null if there is no matched rows.
   */
  def getWithKeyIndex(key: InternalRow): Iterator[ValueRowWithKeyIndex] = {
    throw SparkUnsupportedOperationException()
  }

  /**
   * Returns key index and matched single row.
   * This is for unique key case.
   *
   * Returns null if there is no matched rows.
   */
  def getValueWithKeyIndex(key: InternalRow): ValueRowWithKeyIndex = {
    throw SparkUnsupportedOperationException()
  }

  /**
   * Returns an iterator for keys index and rows of InternalRow type.
   */
  def valuesWithKeyIndex(): Iterator[ValueRowWithKeyIndex] = {
    throw SparkUnsupportedOperationException()
  }

  /**
   * Returns the maximum number of allowed keys index.
   */
  def maxNumKeysIndex: Int = {
    throw SparkUnsupportedOperationException()
  }

  /**
   * Returns true iff all the keys are unique.
   */
  def keyIsUnique: Boolean

  /**
   * Returns an iterator for keys of InternalRow type.
   */
  def keys(): Iterator[InternalRow]

  /**
   * Returns a read-only copy of this, to be safely used in current thread.
   */
  def asReadOnlyCopy(): HashedRelation

  /**
   * Release any used resources.
   */
  def close(): Unit
}

private[execution] object HashedRelation {

  /**
   * Create a HashedRelation from an Iterator of InternalRow.
   *
   * @param allowsNullKey        Allow NULL keys in HashedRelation.
   * @param ignoresDuplicatedKey Ignore rows with duplicated keys in HashedRelation.
   *                             This is only used for semi and anti join without join condition in
   *                             `ShuffledHashJoinExec` only.
   */
  def apply(
      input: Iterator[InternalRow],
      key: Seq[Expression],
      sizeEstimate: Int = 64,
      taskMemoryManager: TaskMemoryManager = null,
      isNullAware: Boolean = false,
      allowsNullKey: Boolean = false,
      ignoresDuplicatedKey: Boolean = false): HashedRelation = {
    val mm = Option(taskMemoryManager).getOrElse {
      new TaskMemoryManager(
        new UnifiedMemoryManager(
          new SparkConf().set(MEMORY_OFFHEAP_ENABLED.key, "false"),
          Runtime.getRuntime.maxMemory,
          Runtime.getRuntime.maxMemory / 2,
          1),
        0)
    }

    if (!input.hasNext && !allowsNullKey) {
      EmptyHashedRelation
    } else if (key.length == 1 && key.head.dataType == LongType && !allowsNullKey) {
      // NOTE: LongHashedRelation does not support NULL keys.
      LongHashedRelation(input, key, sizeEstimate, mm, isNullAware, ignoresDuplicatedKey)
    } else {
      UnsafeHashedRelation(input, key, sizeEstimate, mm, isNullAware, allowsNullKey,
        ignoresDuplicatedKey)
    }
  }
}

/**
 * A wrapper for key index and value in InternalRow type.
 * Designed to be instantiated once per thread and reused.
 */
private[execution] class ValueRowWithKeyIndex {
  private var keyIndex: Int = _
  private var value: InternalRow = _

  /** Updates this ValueRowWithKeyIndex by updating its key index.  Returns itself. */
  def withNewKeyIndex(newKeyIndex: Int): ValueRowWithKeyIndex = {
    keyIndex = newKeyIndex
    this
  }

  /** Updates this ValueRowWithKeyIndex by updating its value.  Returns itself. */
  def withNewValue(newValue: InternalRow): ValueRowWithKeyIndex = {
    value = newValue
    this
  }

  /** Updates this ValueRowWithKeyIndex.  Returns itself. */
  def update(newKeyIndex: Int, newValue: InternalRow): ValueRowWithKeyIndex = {
    keyIndex = newKeyIndex
    value = newValue
    this
  }

  def getKeyIndex: Int = {
    keyIndex
  }

  def getValue: InternalRow = {
    value
  }
}

/**
 * A HashedRelation for UnsafeRow, which is backed BytesToBytesMap.
 *
 * It's serialized in the following format:
 *  [number of keys] [number of fields]
 *  [size of key] [size of value] [key bytes] [bytes for value]
 */
private[joins] class UnsafeHashedRelation(
    private var numKeys: Int,
    private var numFields: Int,
    private var binaryMap: BytesToBytesMap)
  extends HashedRelation with Externalizable with KryoSerializable {

  private[joins] def this() = this(0, 0, null)  // Needed for serialization

  override def keyIsUnique: Boolean = binaryMap.numKeys() == binaryMap.numValues()

  override def asReadOnlyCopy(): UnsafeHashedRelation = {
    new UnsafeHashedRelation(numKeys, numFields, binaryMap)
  }

  override def estimatedSize: Long = binaryMap.getTotalMemoryConsumption

  // re-used in get()/getValue()/getWithKeyIndex()/getValueWithKeyIndex()/valuesWithKeyIndex()
  var resultRow = new UnsafeRow(numFields)

  // re-used in getWithKeyIndex()/getValueWithKeyIndex()/valuesWithKeyIndex()
  val valueRowWithKeyIndex = new ValueRowWithKeyIndex

  override def get(key: InternalRow): Iterator[InternalRow] = {
    val unsafeKey = key.asInstanceOf[UnsafeRow]
    val map = binaryMap  // avoid the compiler error
    val loc = new map.Location  // this could be allocated in stack
    binaryMap.safeLookup(unsafeKey.getBaseObject, unsafeKey.getBaseOffset,
      unsafeKey.getSizeInBytes, loc, unsafeKey.hashCode())
    if (loc.isDefined) {
      new Iterator[UnsafeRow] {
        private var _hasNext = true
        override def hasNext: Boolean = _hasNext
        override def next(): UnsafeRow = {
          resultRow.pointTo(loc.getValueBase, loc.getValueOffset, loc.getValueLength)
          _hasNext = loc.nextValue()
          resultRow
        }
      }
    } else {
      null
    }
  }

  def getValue(key: InternalRow): InternalRow = {
    val unsafeKey = key.asInstanceOf[UnsafeRow]
    val map = binaryMap  // avoid the compiler error
    val loc = new map.Location  // this could be allocated in stack
    binaryMap.safeLookup(unsafeKey.getBaseObject, unsafeKey.getBaseOffset,
      unsafeKey.getSizeInBytes, loc, unsafeKey.hashCode())
    if (loc.isDefined) {
      resultRow.pointTo(loc.getValueBase, loc.getValueOffset, loc.getValueLength)
      resultRow
    } else {
      null
    }
  }

  override def getWithKeyIndex(key: InternalRow): Iterator[ValueRowWithKeyIndex] = {
    val unsafeKey = key.asInstanceOf[UnsafeRow]
    val map = binaryMap  // avoid the compiler error
    val loc = new map.Location  // this could be allocated in stack
    binaryMap.safeLookup(unsafeKey.getBaseObject, unsafeKey.getBaseOffset,
      unsafeKey.getSizeInBytes, loc, unsafeKey.hashCode())
    if (loc.isDefined) {
      valueRowWithKeyIndex.withNewKeyIndex(loc.getKeyIndex)
      new Iterator[ValueRowWithKeyIndex] {
        private var _hasNext = true
        override def hasNext: Boolean = _hasNext
        override def next(): ValueRowWithKeyIndex = {
          resultRow.pointTo(loc.getValueBase, loc.getValueOffset, loc.getValueLength)
          _hasNext = loc.nextValue()
          valueRowWithKeyIndex.withNewValue(resultRow)
        }
      }
    } else {
      null
    }
  }

  override def getValueWithKeyIndex(key: InternalRow): ValueRowWithKeyIndex = {
    val unsafeKey = key.asInstanceOf[UnsafeRow]
    val map = binaryMap  // avoid the compiler error
    val loc = new map.Location  // this could be allocated in stack
    binaryMap.safeLookup(unsafeKey.getBaseObject, unsafeKey.getBaseOffset,
      unsafeKey.getSizeInBytes, loc, unsafeKey.hashCode())
    if (loc.isDefined) {
      resultRow.pointTo(loc.getValueBase, loc.getValueOffset, loc.getValueLength)
      valueRowWithKeyIndex.update(loc.getKeyIndex, resultRow)
    } else {
      null
    }
  }

  override def valuesWithKeyIndex(): Iterator[ValueRowWithKeyIndex] = {
    val iter = binaryMap.iteratorWithKeyIndex()

    new Iterator[ValueRowWithKeyIndex] {
      override def hasNext: Boolean = iter.hasNext

      override def next(): ValueRowWithKeyIndex = {
        if (!hasNext) {
          throw QueryExecutionErrors.endOfIteratorError()
        }
        val loc = iter.next()
        resultRow.pointTo(loc.getValueBase, loc.getValueOffset, loc.getValueLength)
        valueRowWithKeyIndex.update(loc.getKeyIndex, resultRow)
      }
    }
  }

  override def maxNumKeysIndex: Int = {
    binaryMap.maxNumKeysIndex
  }

  override def keys(): Iterator[InternalRow] = {
    val iter = binaryMap.iterator()

    new Iterator[InternalRow] {
      val unsafeRow = new UnsafeRow(numKeys)

      override def hasNext: Boolean = {
        iter.hasNext
      }

      override def next(): InternalRow = {
        if (!hasNext) {
          throw QueryExecutionErrors.endOfIteratorError()
        } else {
          val loc = iter.next()
          unsafeRow.pointTo(loc.getKeyBase, loc.getKeyOffset, loc.getKeyLength)
          unsafeRow
        }
      }
    }
  }

  override def close(): Unit = {
    binaryMap.free()
  }

  override def writeExternal(out: ObjectOutput): Unit = Utils.tryOrIOException {
    write(out.writeInt, out.writeLong, out.write)
  }

  override def write(kryo: Kryo, out: Output): Unit = Utils.tryOrIOException {
    write(out.writeInt, out.writeLong, out.write)
  }

  private def write(
      writeInt: (Int) => Unit,
      writeLong: (Long) => Unit,
      writeBuffer: (Array[Byte], Int, Int) => Unit) : Unit = {
    writeInt(numKeys)
    writeInt(numFields)
    // TODO: move these into BytesToBytesMap
    writeLong(binaryMap.numKeys())
    writeLong(binaryMap.numValues())

    var buffer = new Array[Byte](64)
    def write(base: Object, offset: Long, length: Int): Unit = {
      if (buffer.length < length) {
        buffer = new Array[Byte](length)
      }
      Platform.copyMemory(base, offset, buffer, Platform.BYTE_ARRAY_OFFSET, length)
      writeBuffer(buffer, 0, length)
    }

    val iter = binaryMap.iterator()
    while (iter.hasNext) {
      val loc = iter.next()
      // [key size] [values size] [key bytes] [value bytes]
      writeInt(loc.getKeyLength)
      writeInt(loc.getValueLength)
      write(loc.getKeyBase, loc.getKeyOffset, loc.getKeyLength)
      write(loc.getValueBase, loc.getValueOffset, loc.getValueLength)
    }
  }

  override def readExternal(in: ObjectInput): Unit = Utils.tryOrIOException {
    read(() => in.readInt(), () => in.readLong(), in.readFully)
  }

  private def read(
      readInt: () => Int,
      readLong: () => Long,
      readBuffer: (Array[Byte], Int, Int) => Unit): Unit = {
    numKeys = readInt()
    numFields = readInt()
    resultRow = new UnsafeRow(numFields)
    val nKeys = readLong()
    val nValues = readLong()
    // This is used in Broadcast, shared by multiple tasks, so we use on-heap memory
    // TODO(josh): This needs to be revisited before we merge this patch; making this change now
    // so that tests compile:
    val taskMemoryManager = new TaskMemoryManager(
      new UnifiedMemoryManager(
        new SparkConf().set(MEMORY_OFFHEAP_ENABLED.key, "false"),
        Runtime.getRuntime.maxMemory,
        Runtime.getRuntime.maxMemory / 2,
        1),
      0)

    val pageSizeBytes = Option(SparkEnv.get).map(_.memoryManager.pageSizeBytes)
      .getOrElse(new SparkConf().get(BUFFER_PAGESIZE).getOrElse(16L * 1024 * 1024))

    // TODO(josh): We won't need this dummy memory manager after future refactorings; revisit
    // during code review

    binaryMap = new BytesToBytesMap(
      taskMemoryManager,
      (nKeys * 1.5 + 1).toInt, // reduce hash collision
      pageSizeBytes)

    var i = 0
    var keyBuffer = new Array[Byte](1024)
    var valuesBuffer = new Array[Byte](1024)
    while (i < nValues) {
      val keySize = readInt()
      val valuesSize = readInt()
      if (keySize > keyBuffer.length) {
        keyBuffer = new Array[Byte](keySize)
      }
      readBuffer(keyBuffer, 0, keySize)
      if (valuesSize > valuesBuffer.length) {
        valuesBuffer = new Array[Byte](valuesSize)
      }
      readBuffer(valuesBuffer, 0, valuesSize)

      val loc = binaryMap.lookup(keyBuffer, Platform.BYTE_ARRAY_OFFSET, keySize)
      val putSucceeded = loc.append(keyBuffer, Platform.BYTE_ARRAY_OFFSET, keySize,
        valuesBuffer, Platform.BYTE_ARRAY_OFFSET, valuesSize)
      if (!putSucceeded) {
        binaryMap.free()
        throw QueryExecutionErrors.cannotAllocateMemoryToGrowBytesToBytesMapError()
      }
      i += 1
    }
  }

  override def read(kryo: Kryo, in: Input): Unit = Utils.tryOrIOException {
    read(() => in.readInt(), () => in.readLong(), in.readBytes)
  }
}

private[joins] object UnsafeHashedRelation {

  def apply(
      input: Iterator[InternalRow],
      key: Seq[Expression],
      sizeEstimate: Int,
      taskMemoryManager: TaskMemoryManager,
      isNullAware: Boolean = false,
      allowsNullKey: Boolean = false,
      ignoresDuplicatedKey: Boolean = false): HashedRelation = {
    require(!(isNullAware && allowsNullKey),
      "isNullAware and allowsNullKey cannot be enabled at same time")

    val pageSizeBytes = Option(SparkEnv.get).map(_.memoryManager.pageSizeBytes)
      .getOrElse(new SparkConf().get(BUFFER_PAGESIZE).getOrElse(16L * 1024 * 1024))
    val binaryMap = new BytesToBytesMap(
      taskMemoryManager,
      // Only 70% of the slots can be used before growing, more capacity help to reduce collision
      (sizeEstimate * 1.5 + 1).toInt,
      pageSizeBytes)

    // Create a mapping of buildKeys -> rows
    val keyGenerator = UnsafeProjection.create(key)
    var numFields = 0
    while (input.hasNext) {
      val row = input.next().asInstanceOf[UnsafeRow]
      numFields = row.numFields()
      val key = keyGenerator(row)
      if (!key.anyNull || allowsNullKey) {
        val loc = binaryMap.lookup(key.getBaseObject, key.getBaseOffset, key.getSizeInBytes)
        if (!(ignoresDuplicatedKey && loc.isDefined)) {
          val success = loc.append(
            key.getBaseObject, key.getBaseOffset, key.getSizeInBytes,
            row.getBaseObject, row.getBaseOffset, row.getSizeInBytes)
          if (!success) {
            binaryMap.free()
            throw QueryExecutionErrors.cannotAcquireMemoryToBuildUnsafeHashedRelationError()
          }
        }
      } else if (isNullAware) {
        binaryMap.free()
        return HashedRelationWithAllNullKeys
      }
    }

    new UnsafeHashedRelation(key.size, numFields, binaryMap)
  }
}

/**
 * An append-only hash map mapping from key of Long to UnsafeRow.
 *
 * The underlying bytes of all values (UnsafeRows) are packed together as a single byte array
 * (`page`) in this format:
 *
 *  [bytes of row1][address1][bytes of row2][address1] ...
 *
 *  address1 (8 bytes) is the offset and size of next value for the same key as row1, any key
 *  could have multiple values. the address at the end of last value for every key is 0.
 *
 * The keys and addresses of their values could be stored in two modes:
 *
 * 1) sparse mode: the keys and addresses are stored in `array` as:
 *
 *  [key1][address1][key2][address2]...[]
 *
 *  address1 (Long) is the offset (in `page`) and size of the value for key1. The position of key1
 *  is determined by `key1 % cap`. Quadratic probing with triangular numbers is used to address
 *  hash collision.
 *
 * 2) dense mode: all the addresses are packed into a single array of long, as:
 *
 *  [address1] [address2] ...
 *
 *  address1 (Long) is the offset (in `page`) and size of the value for key1, the position is
 *  determined by `key1 - minKey`.
 *
 * The map is created as sparse mode, then key-value could be appended into it. Once finish
 * appending, caller could call optimize() to try to turn the map into dense mode, which is faster
 * to probe.
 *
 * see http://java-performance.info/implementing-world-fastest-java-int-to-int-hash-map/
 */
private[execution] final class LongToUnsafeRowMap(
    val mm: TaskMemoryManager,
    capacity: Int,
    ignoresDuplicatedKey: Boolean = false)
  extends MemoryConsumer(mm, MemoryMode.ON_HEAP) with Externalizable with KryoSerializable {

  // Whether the keys are stored in dense mode or not.
  private var isDense = false

  // The minimum key
  private var minKey = Long.MaxValue

  // The maximum key
  private var maxKey = Long.MinValue

  // The array to store the key and offset of UnsafeRow in the page.
  //
  // Sparse mode: [key1] [offset1 | size1] [key2] [offset | size2] ...
  // Dense mode: [offset1 | size1] [offset2 | size2]
  private var array: Array[Long] = null
  private var mask: Int = 0

  // The page to store all bytes of UnsafeRow and the pointer to next rows.
  // [row1][pointer1] [row2][pointer2]
  private var page: Array[Long] = null

  // Current write cursor in the page.
  private var cursor: Long = Platform.LONG_ARRAY_OFFSET

  // The number of bits for size in address
  private val SIZE_BITS = 28
  private val SIZE_MASK = 0xfffffff

  // The total number of values of all keys.
  private var numValues = 0L

  // The number of unique keys.
  private var numKeys = 0L

  // needed by serializer
  def this() = {
    this(
      new TaskMemoryManager(
        new UnifiedMemoryManager(
          new SparkConf().set(MEMORY_OFFHEAP_ENABLED.key, "false"),
          Runtime.getRuntime.maxMemory,
          Runtime.getRuntime.maxMemory / 2,
          1),
        0),
      0)
  }

  private def ensureAcquireMemory(size: Long): Unit = {
    // do not support spilling
    val got = acquireMemory(size)
    if (got < size) {
      freeMemory(got)
      throw QueryExecutionErrors.cannotAcquireMemoryToBuildLongHashedRelationError(size, got)
    }
  }

  private def init(): Unit = {
    if (mm != null) {
      require(capacity < 512000000, "Cannot broadcast 512 million or more rows")
      var n = 1
      while (n < capacity) n *= 2
      ensureAcquireMemory(n * 2L * 8 + (1 << 20))
      array = new Array[Long](n * 2)
      mask = n * 2 - 2
      page = new Array[Long](1 << 17)  // 1M bytes
    }
  }

  init()

  def spill(size: Long, trigger: MemoryConsumer): Long = 0L

  /**
   * Returns whether all the keys are unique.
   */
  def keyIsUnique: Boolean = numKeys == numValues

  /**
   * Returns total memory consumption.
   */
  def getTotalMemoryConsumption: Long = array.length * 8L + page.length * 8L

  /**
   * Returns the first slot of array that store the keys (sparse mode).
   */
  private def firstSlot(key: Long): Int = {
    val h = key * 0x9E3779B9L
    (h ^ (h >> 32)).toInt & mask
  }

  /**
   * Returns the next probe in the array.
   */
  private def nextSlot(pos: Int): Int = (pos + 2) & mask

  private[this] def toAddress(offset: Long, size: Int): Long = {
    ((offset - Platform.LONG_ARRAY_OFFSET) << SIZE_BITS) | size
  }

  private[this] def toOffset(address: Long): Long = {
    (address >>> SIZE_BITS) + Platform.LONG_ARRAY_OFFSET
  }

  private[this] def toSize(address: Long): Int = {
    (address & SIZE_MASK).toInt
  }

  private def getRow(address: Long, resultRow: UnsafeRow): UnsafeRow = {
    resultRow.pointTo(page, toOffset(address), toSize(address))
    resultRow
  }

  /**
   * Returns the single UnsafeRow for given key, or null if not found.
   */
  def getValue(key: Long, resultRow: UnsafeRow): UnsafeRow = {
    if (isDense) {
      if (key >= minKey && key <= maxKey) {
        val value = array((key - minKey).toInt)
        if (value > 0) {
          return getRow(value, resultRow)
        }
      }
    } else {
      var pos = firstSlot(key)
      while (array(pos + 1) != 0) {
        if (array(pos) == key) {
          return getRow(array(pos + 1), resultRow)
        }
        pos = nextSlot(pos)
      }
    }
    null
  }

  /**
   * Returns an iterator of UnsafeRow for multiple linked values.
   */
  private def valueIter(address: Long, resultRow: UnsafeRow): Iterator[UnsafeRow] = {
    new Iterator[UnsafeRow] {
      var addr = address
      override def hasNext: Boolean = addr != 0
      override def next(): UnsafeRow = {
        val offset = toOffset(addr)
        val size = toSize(addr)
        resultRow.pointTo(page, offset, size)
        addr = Platform.getLong(page, offset + size)
        resultRow
      }
    }
  }

  /**
   * Returns an iterator for all the values for the given key, or null if no value found.
   */
  def get(key: Long, resultRow: UnsafeRow): Iterator[UnsafeRow] = {
    if (isDense) {
      if (key >= minKey && key <= maxKey) {
        val value = array((key - minKey).toInt)
        if (value > 0) {
          return valueIter(value, resultRow)
        }
      }
    } else {
      var pos = firstSlot(key)
      while (array(pos + 1) != 0) {
        if (array(pos) == key) {
          return valueIter(array(pos + 1), resultRow)
        }
        pos = nextSlot(pos)
      }
    }
    null
  }

  /**
   * Builds an iterator on a sparse array.
   */
  def keys(): Iterator[InternalRow] = {
    val row = new GenericInternalRow(1)
    // a) in dense mode the array stores the address
    //  => (k, v) = (minKey + index, array(index))
    // b) in sparse mode the array stores both the key and the address
    //  => (k, v) = (array(index), array(index+1))
    new Iterator[InternalRow] {
      // cursor that indicates the position of the next key which was not read by a next() call
      var pos = 0
      // when we iterate in dense mode we need to jump two positions at a time
      val step = if (isDense) 0 else 1

      override def hasNext: Boolean = {
        // go to the next key if the current key slot is empty
        while (pos + step < array.length) {
          if (array(pos + step) > 0) {
            return true
          }
          pos += step + 1
        }
        false
      }

      override def next(): InternalRow = {
        if (!hasNext) {
          throw QueryExecutionErrors.endOfIteratorError()
        } else {
          // the key is retrieved based on the map mode
          val ret = if (isDense) minKey + pos else array(pos)
          // advance the cursor to the next index
          pos += step + 1
          row.setLong(0, ret)
          row
        }
      }
    }
  }

  /**
   * Appends the key and row into this map.
   */
  def append(key: Long, row: UnsafeRow): Unit = {
    val sizeInBytes = row.getSizeInBytes
    if (sizeInBytes >= (1 << SIZE_BITS)) {
      throw QueryExecutionErrors.rowLargerThan256MUnsupportedError()
    }

    val pos = findKeyPosition(key)
    if (ignoresDuplicatedKey && array(pos + 1) != 0) {
      return
    }

    if (key < minKey) {
      minKey = key
    }
    if (key > maxKey) {
      maxKey = key
    }

    grow(row.getSizeInBytes)

    // copy the bytes of UnsafeRow
    val offset = cursor
    Platform.copyMemory(row.getBaseObject, row.getBaseOffset, page, cursor, row.getSizeInBytes)
    cursor += row.getSizeInBytes
    Platform.putLong(page, cursor, 0)
    cursor += 8
    numValues += 1
    updateIndex(key, pos, toAddress(offset, row.getSizeInBytes))
  }

  private def findKeyPosition(key: Long): Int = {
    var pos = firstSlot(key)
    assert(numKeys < array.length / 2)
    while (array(pos) != key && array(pos + 1) != 0) {
      pos = nextSlot(pos)
    }
    pos
  }

  /**
   * Update the address in array for given key.
   */
  private def updateIndex(key: Long, pos: Int, address: Long): Unit = {
    if (array(pos + 1) == 0) {
      // this is the first value for this key, put the address in array.
      array(pos) = key
      array(pos + 1) = address
      numKeys += 1
      if (numKeys * 4 > array.length) {
        // reach half of the capacity
        if (array.length < (1 << 30)) {
          // Cannot allocate an array with 2G elements
          growArray()
        } else if (numKeys > array.length / 2 * 0.75) {
          // The fill ratio should be less than 0.75
          throw QueryExecutionErrors.cannotBuildHashedRelationWithUniqueKeysExceededError()
        }
      }
    } else {
      // there are some values for this key, put the address in the front of them.
      val pointer = toOffset(address) + toSize(address)
      Platform.putLong(page, pointer, array(pos + 1))
      array(pos + 1) = address
    }
  }

  private def grow(inputRowSize: Int): Unit = {
    // There is 8 bytes for the pointer to next value
    val neededNumWords = (cursor - Platform.LONG_ARRAY_OFFSET + 8 + inputRowSize + 7) / 8
    if (neededNumWords > page.length) {
      if (neededNumWords > (1 << 30)) {
        throw QueryExecutionErrors.cannotBuildHashedRelationLargerThan8GError()
      }
      val newNumWords = math.max(neededNumWords, math.min(page.length * 2, 1 << 30))
      ensureAcquireMemory(newNumWords * 8L)
      val newPage = new Array[Long](newNumWords.toInt)
      Platform.copyMemory(page, Platform.LONG_ARRAY_OFFSET, newPage, Platform.LONG_ARRAY_OFFSET,
        cursor - Platform.LONG_ARRAY_OFFSET)
      val used = page.length
      page = newPage
      freeMemory(used * 8L)
    }
  }

  private def growArray(): Unit = {
    var old_array = array
    val n = array.length
    numKeys = 0
    ensureAcquireMemory(n * 2 * 8L)
    array = new Array[Long](n * 2)
    mask = n * 2 - 2
    var i = 0
    while (i < old_array.length) {
      if (old_array(i + 1) > 0) {
        val key = old_array(i)
        updateIndex(key, findKeyPosition(key), old_array(i + 1))
      }
      i += 2
    }
    old_array = null  // release the reference to old array
    freeMemory(n * 8L)
  }

  /**
   * Try to turn the map into dense mode, which is faster to probe.
   */
  def optimize(): Unit = {
    val range = maxKey - minKey
    // Convert to dense mode if it does not require more memory or could fit within L1 cache
    // SPARK-16740: Make sure range doesn't overflow if minKey has a large negative value
    if (range >= 0 && (range < array.length || range < 1024)) {
      try {
        ensureAcquireMemory((range + 1) * 8L)
      } catch {
        case e: SparkException =>
          // there is no enough memory to convert
          return
      }
      val denseArray = new Array[Long]((range + 1).toInt)
      var i = 0
      while (i < array.length) {
        if (array(i + 1) > 0) {
          val idx = (array(i) - minKey).toInt
          denseArray(idx) = array(i + 1)
        }
        i += 2
      }
      val old_length = array.length
      array = denseArray
      isDense = true
      freeMemory(old_length * 8L)
    }
  }

  /**
   * Free all the memory acquired by this map.
   */
  def free(): Unit = {
    if (page != null) {
      freeMemory(page.length * 8L)
      page = null
    }
    if (array != null) {
      freeMemory(array.length * 8L)
      array = null
    }
  }

  private def writeLongArray(
      writeBuffer: (Array[Byte], Int, Int) => Unit,
      arr: Array[Long],
      len: Int): Unit = {
    val buffer = new Array[Byte](4 << 10)
    var offset: Long = Platform.LONG_ARRAY_OFFSET
    val end = len * 8L + Platform.LONG_ARRAY_OFFSET
    while (offset < end) {
      val size = Math.min(buffer.length, end - offset)
      Platform.copyMemory(arr, offset, buffer, Platform.BYTE_ARRAY_OFFSET, size)
      writeBuffer(buffer, 0, size.toInt)
      offset += size
    }
  }

  private def write(
      writeBoolean: (Boolean) => Unit,
      writeLong: (Long) => Unit,
      writeBuffer: (Array[Byte], Int, Int) => Unit): Unit = {
    writeBoolean(isDense)
    writeLong(minKey)
    writeLong(maxKey)
    writeLong(numKeys)
    writeLong(numValues)

    writeLong(array.length)
    writeLongArray(writeBuffer, array, array.length)
    val used = ((cursor - Platform.LONG_ARRAY_OFFSET) / 8).toInt
    writeLong(used)
    writeLongArray(writeBuffer, page, used)
  }

  override def writeExternal(output: ObjectOutput): Unit = {
    write(output.writeBoolean, output.writeLong, output.write)
  }

  override def write(kryo: Kryo, out: Output): Unit = {
    write(out.writeBoolean, out.writeLong, out.write)
  }

  private def readLongArray(
      readBuffer: (Array[Byte], Int, Int) => Unit,
      length: Int): Array[Long] = {
    val array = new Array[Long](length)
    val buffer = new Array[Byte](4 << 10)
    var offset: Long = Platform.LONG_ARRAY_OFFSET
    val end = length * 8L + Platform.LONG_ARRAY_OFFSET
    while (offset < end) {
      val size = Math.min(buffer.length, end - offset)
      readBuffer(buffer, 0, size.toInt)
      Platform.copyMemory(buffer, Platform.BYTE_ARRAY_OFFSET, array, offset, size)
      offset += size
    }
    array
  }

  private def read(
      readBoolean: () => Boolean,
      readLong: () => Long,
      readBuffer: (Array[Byte], Int, Int) => Unit): Unit = {
    isDense = readBoolean()
    minKey = readLong()
    maxKey = readLong()
    numKeys = readLong()
    numValues = readLong()

    val length = readLong().toInt
    mask = length - 2
    array = readLongArray(readBuffer, length)
    val pageLength = readLong().toInt
    page = readLongArray(readBuffer, pageLength)
    // Restore cursor variable to make this map able to be serialized again on executors.
    cursor = pageLength * 8 + Platform.LONG_ARRAY_OFFSET
  }

  override def readExternal(in: ObjectInput): Unit = {
    read(() => in.readBoolean(), () => in.readLong(), in.readFully)
  }

  override def read(kryo: Kryo, in: Input): Unit = {
    read(() => in.readBoolean(), () => in.readLong(), in.readBytes)
  }
}

class LongHashedRelation(
    private var nFields: Int,
    private var map: LongToUnsafeRowMap) extends HashedRelation with Externalizable {

  private var resultRow: UnsafeRow = new UnsafeRow(nFields)

  // Needed for serialization (it is public to make Java serialization work)
  def this() = this(0, null)

  override def asReadOnlyCopy(): LongHashedRelation = new LongHashedRelation(nFields, map)

  override def estimatedSize: Long = map.getTotalMemoryConsumption

  override def get(key: InternalRow): Iterator[InternalRow] = {
    if (key.isNullAt(0)) {
      null
    } else {
      get(key.getLong(0))
    }
  }

  override def getValue(key: InternalRow): InternalRow = {
    if (key.isNullAt(0)) {
      null
    } else {
      getValue(key.getLong(0))
    }
  }

  override def get(key: Long): Iterator[InternalRow] = map.get(key, resultRow)

  override def getValue(key: Long): InternalRow = map.getValue(key, resultRow)

  override def keyIsUnique: Boolean = map.keyIsUnique

  override def close(): Unit = {
    map.free()
  }

  override def writeExternal(out: ObjectOutput): Unit = {
    out.writeInt(nFields)
    out.writeObject(map)
  }

  override def readExternal(in: ObjectInput): Unit = {
    nFields = in.readInt()
    resultRow = new UnsafeRow(nFields)
    map = in.readObject().asInstanceOf[LongToUnsafeRowMap]
  }

  /**
   * Returns an iterator for keys of InternalRow type.
   */
  override def keys(): Iterator[InternalRow] = map.keys()

  override def getWithKeyIndex(key: InternalRow): Iterator[ValueRowWithKeyIndex] = {
    throw SparkUnsupportedOperationException()
  }

  override def getValueWithKeyIndex(key: InternalRow): ValueRowWithKeyIndex = {
    throw SparkUnsupportedOperationException()
  }

  override def valuesWithKeyIndex(): Iterator[ValueRowWithKeyIndex] = {
    throw SparkUnsupportedOperationException()
  }

  override def maxNumKeysIndex: Int = {
    throw SparkUnsupportedOperationException()
  }
}

/**
 * Create hashed relation with key that is long.
 */
private[joins] object LongHashedRelation {
  def apply(
      input: Iterator[InternalRow],
      key: Seq[Expression],
      sizeEstimate: Int,
      taskMemoryManager: TaskMemoryManager,
      isNullAware: Boolean = false,
      ignoresDuplicatedKey: Boolean = false): HashedRelation = {

    val map = new LongToUnsafeRowMap(taskMemoryManager, sizeEstimate, ignoresDuplicatedKey)
    val keyGenerator = UnsafeProjection.create(key)

    // Create a mapping of key -> rows
    var numFields = 0
    while (input.hasNext) {
      val unsafeRow = input.next().asInstanceOf[UnsafeRow]
      numFields = unsafeRow.numFields()
      val rowKey = keyGenerator(unsafeRow)
      if (!rowKey.isNullAt(0)) {
        val key = rowKey.getLong(0)
        map.append(key, unsafeRow)
      } else if (isNullAware) {
        map.free()
        return HashedRelationWithAllNullKeys
      }
    }
    map.optimize()
    new LongHashedRelation(numFields, map)
  }
}

/**
 * A special HashedRelation indicating that it's built from a empty input:Iterator[InternalRow].
 * get & getValue will return null just like
 * empty LongHashedRelation or empty UnsafeHashedRelation does.
 */
case object EmptyHashedRelation extends HashedRelation {
  override def get(key: Long): Iterator[InternalRow] = null

  override def get(key: InternalRow): Iterator[InternalRow] = null

  override def getValue(key: Long): InternalRow = null

  override def getValue(key: InternalRow): InternalRow = null

  override def asReadOnlyCopy(): EmptyHashedRelation.type = this

  override def keyIsUnique: Boolean = true

  override def keys(): Iterator[InternalRow] = {
    Iterator.empty
  }

  override def close(): Unit = {}

  override def estimatedSize: Long = 0
}

/**
 * A special HashedRelation indicating that it's built from a non-empty input:Iterator[InternalRow]
 * with all the keys to be null.
 */
case object HashedRelationWithAllNullKeys extends HashedRelation {
  override def get(key: InternalRow): Iterator[InternalRow] = {
    throw SparkUnsupportedOperationException()
  }

  override def getValue(key: InternalRow): InternalRow = {
    throw SparkUnsupportedOperationException()
  }

  override def asReadOnlyCopy(): HashedRelationWithAllNullKeys.type = this

  override def keyIsUnique: Boolean = true

  override def keys(): Iterator[InternalRow] = {
    throw SparkUnsupportedOperationException()
  }

  override def close(): Unit = {}

  override def estimatedSize: Long = 0
}

/** The HashedRelationBroadcastMode requires that rows are broadcasted as a HashedRelation. */
case class HashedRelationBroadcastMode(key: Seq[Expression], isNullAware: Boolean = false)
  extends BroadcastMode {

  override def transform(rows: Array[InternalRow]): HashedRelation = {
    transform(rows.iterator, Some(rows.length))
  }

  override def transform(
      rows: Iterator[InternalRow],
      sizeHint: Option[Long]): HashedRelation = {
    sizeHint match {
      case Some(numRows) =>
        HashedRelation(rows, key, numRows.toInt, isNullAware = isNullAware)
      case None =>
        HashedRelation(rows, key, isNullAware = isNullAware)
    }
  }

  override lazy val canonicalized: HashedRelationBroadcastMode = {
    this.copy(key = key.map(_.canonicalized))
  }
}
