/*
 * 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.catalyst.expressions

import scala.util.{Either, Left, Right}

import org.apache.spark.{QueryContext, SparkException}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.analysis._
import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, CodeGenerator, ExprCode}
import org.apache.spark.sql.catalyst.trees.TreePattern.{EXTRACT_VALUE, TreePattern}
import org.apache.spark.sql.catalyst.util.{quoteIdentifier, ArrayData, GenericArrayData, MapData, TypeUtils}
import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryErrorsBase, QueryExecutionErrors}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._

////////////////////////////////////////////////////////////////////////////////////////////////////
// This file defines all the expressions to extract values out of complex types.
// For example, getting a field out of an array, map, or struct.
////////////////////////////////////////////////////////////////////////////////////////////////////


object ExtractValue {
  /**
   * Returns the resolved `ExtractValue` using the `extractValue` method call. In case the method
   * returns `None`, it throws.
   *
   * See `extractValue` doc for more info.
   */
  def apply(
      child: Expression,
      extraction: Expression,
      resolver: Resolver): Expression = {
    extractValue(child, extraction, resolver) match {
      case Left(expression) => expression
      case Right(throwable) => throw throwable
    }
  }

  /**
   * Resolution-time variant of [[apply]]: extracting a field/element/key from a NULL (`NullType`)
   * base yields NULL (SQL NULL propagation) instead of throwing `INVALID_EXTRACT_BASE_FIELD_TYPE`.
   * A `NullType` column can arise e.g. from schema evolution with missing columns. This is used by
   * the user-facing extraction resolution sites (multipart name resolution and
   * `UnresolvedExtractValue` resolution). `extractValue` itself is left unchanged, so the other
   * direct consumers keep their prior (throwing) behavior.
   */
  def applyOrNull(
      child: Expression,
      extraction: Expression,
      resolver: Resolver): Expression = {
    if (child.dataType == NullType) {
      Literal(null, NullType)
    } else {
      apply(child, extraction, resolver)
    }
  }

  /**
   * Returns the resolved `ExtractValue`. It will return one kind of concrete `ExtractValue`,
   * depend on the type of `child` and `extraction`.
   *
   *   `child`      |    `extraction`    |    concrete `ExtractValue`
   * ----------------------------------------------------------------
   *    Struct      |   Literal String   |        GetStructField
   * Array[Struct]  |   Literal String   |     GetArrayStructFields
   *    Array       |   Integral type    |         GetArrayItem
   *     Map        |   map key type     |         GetMapValue
   */
  def extractValue(
      child: Expression,
      extraction: Expression,
      resolver: Resolver): Either[Expression, Throwable] = {
    (child.dataType, extraction) match {
      case (StructType(fields), NonNullLiteral(v, _: StringType)) =>
        val fieldName = v.toString
        findField(fields, fieldName, resolver) match {
          case Left(ordinal) => Left(GetStructField(child, ordinal, Some(fieldName)))
          case Right(throwable) => Right(throwable)
        }

      case (ArrayType(StructType(fields), containsNull), NonNullLiteral(v, _: StringType)) =>
        val fieldName = v.toString
        findField(fields, fieldName, resolver) match {
          case Left(ordinal) =>
            Left(
              GetArrayStructFields(
                child,
                fields(ordinal).copy(name = fieldName),
                ordinal,
                fields.length,
                containsNull || fields(ordinal).nullable
              )
            )
          case Right(throwable) => Right(throwable)
        }

      case (_: ArrayType, _) => Left(GetArrayItem(child, extraction))

      case (MapType(_, _, _), _) => Left(GetMapValue(child, extraction))

      case (otherType, _) => Right(
        QueryCompilationErrors.dataTypeUnsupportedByExtractValueError(
          child.dataType,
          extraction,
          child
        )
      )
    }
  }

  /**
   * Check that `attribute` can be fully extracted using the given nestedFields and `extractorKey`.
   * This method first tries to extract the innermost nested field using `nestedFields`. If such
   * nested field is found, try to extract the value using `extractorKey`, if it exists. This
   * method returns true if all attempted extractions are successful.
   */
  def isExtractable(
      attribute: Attribute,
      nestedFields: Seq[String],
      extractorKey: Option[Expression] = None,
      resolver: Resolver): Boolean = {
    val withExtractedNestedFields = nestedFields
      .foldLeft(Some(attribute): Option[Expression]) {
        case (Some(expression), field) =>
          // Extraction from a NULL (NullType) base propagates NULL rather than failing, matching
          // the user-facing resolution sites (which use applyOrNull). Treating it as extractable
          // here keeps the NullType candidate in single-pass NameScope candidate filtering so it
          // resolves consistently with the legacy analyzer.
          if (expression.dataType == NullType) {
            Some(Literal(null, NullType))
          } else {
            ExtractValue.extractValue(
              child = expression,
              extraction = Literal(field),
              resolver = resolver
            ) match {
              case Left(e) => Some(e)
              case Right(_) => None
            }
          }
        case _ =>
          None
      }

    withExtractedNestedFields match {
      case Some(nestedField) if extractorKey.isDefined =>
        ExtractValue.extractValue(
          child = nestedField,
          extraction = extractorKey.get,
          resolver = resolver
        ) match {
          case Left(e) => true
          case Right(_) => false
        }
      case Some(_) => true
      case None => false
    }
  }

  /**
   * Find the ordinal of StructField, report error if no desired field or over one
   * desired fields are found.
   */
  private def findField(
      fields: Array[StructField], fieldName: String, resolver: Resolver): Either[Int, Throwable] = {
    val checkField = (f: StructField) => resolver(f.name, fieldName)
    val ordinal = fields.indexWhere(checkField)
    if (ordinal == -1) {
      Right(QueryCompilationErrors.noSuchStructFieldInGivenFieldsError(fieldName, fields))
    } else if (fields.indexWhere(checkField, ordinal + 1) != -1) {
      val numberOfAppearance = fields.count(checkField)
      Right(QueryCompilationErrors.ambiguousReferenceToFieldsError(fieldName, numberOfAppearance))
    } else {
      Left(ordinal)
    }
  }
}

trait ExtractValue extends Expression with QueryErrorsBase {
  override def nullIntolerant: Boolean = true
  final override val nodePatterns: Seq[TreePattern] = Seq(EXTRACT_VALUE)
  val child: Expression
}

/**
 * Returns the value of fields in the Struct `child`.
 *
 * No need to do type checking since it is handled by [[ExtractValue]].
 *
 * Note that we can pass in the field name directly to keep case preserving in `toString`.
 * For example, when get field `yEAr` from `<year: int, month: int>`, we should pass in `yEAr`.
 */
case class GetStructField(child: Expression, ordinal: Int, name: Option[String] = None)
  extends UnaryExpression with ExtractValue with ExpectsInputTypes {

  override def inputTypes: Seq[AbstractDataType] = Seq(StructType, IntegralType)

  lazy val childSchema = child.dataType match {
    case st: StructType => st
    case other =>
      throw SparkException.internalError(
        s"GetStructField requires a StructType child, but got ${other.catalogString}. " +
        s"The child type may have been changed by a plan transformation.")
  }

  override lazy val canonicalized: Expression = {
    copy(child = child.canonicalized, name = None)
  }

  override def dataType: DataType = childSchema(ordinal).dataType
  override def nullable: Boolean = child.nullable || childSchema(ordinal).nullable

  override def toString: String = {
    val fieldName = if (resolved) childSchema(ordinal).name else s"_$ordinal"
    s"$child.${name.getOrElse(fieldName)}"
  }

  def extractFieldName: String = name.getOrElse(childSchema(ordinal).name)

  override def sql: String =
    child.sql + s".${quoteIdentifier(extractFieldName)}"

  protected override def nullSafeEval(input: Any): Any =
    input.asInstanceOf[InternalRow].get(ordinal, childSchema(ordinal).dataType)

  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
    nullSafeCodeGen(ctx, ev, eval => {
      if (nullable) {
        s"""
          if ($eval.isNullAt($ordinal)) {
            ${ev.isNull} = true;
          } else {
            ${ev.value} = ${CodeGenerator.getValue(eval, dataType, ordinal.toString)};
          }
        """
      } else {
        s"""
          ${ev.value} = ${CodeGenerator.getValue(eval, dataType, ordinal.toString)};
        """
      }
    })
  }

  override protected def withNewChildInternal(newChild: Expression): GetStructField =
    copy(child = newChild)

  def metadata: Metadata = childSchema(ordinal).metadata
}

/**
 * For a child whose data type is an array of structs, extracts the `ordinal`-th fields of all array
 * elements, and returns them as a new array.
 *
 * No need to do type checking since it is handled by [[ExtractValue]].
 */
case class GetArrayStructFields(
    child: Expression,
    field: StructField,
    ordinal: Int,
    numFields: Int,
    containsNull: Boolean) extends UnaryExpression with ExtractValue {

  override def checkInputDataTypes(): TypeCheckResult = child.dataType match {
    case ArrayType(_: StructType, _) => TypeCheckResult.TypeCheckSuccess
    // This should never happen, unless we hit a bug.
    case other => TypeCheckResult.TypeCheckFailure(
      "GetArrayStructFields.child must be array of struct type, but got " + other)
  }

  override def dataType: DataType = ArrayType(field.dataType, containsNull)
  override def toString: String = s"$child.${field.name}"
  override def sql: String = s"${child.sql}.${quoteIdentifier(field.name)}"

  protected override def nullSafeEval(input: Any): Any = {
    val array = input.asInstanceOf[ArrayData]
    val length = array.numElements()
    val result = new Array[Any](length)
    var i = 0
    while (i < length) {
      if (array.isNullAt(i)) {
        result(i) = null
      } else {
        val row = array.getStruct(i, numFields)
        if (row.isNullAt(ordinal)) {
          result(i) = null
        } else {
          result(i) = row.get(ordinal, field.dataType)
        }
      }
      i += 1
    }
    new GenericArrayData(result)
  }

  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
    val arrayClass = classOf[GenericArrayData].getName
    nullSafeCodeGen(ctx, ev, eval => {
      val n = ctx.freshName("n")
      val values = ctx.freshName("values")
      val j = ctx.freshName("j")
      val row = ctx.freshName("row")
      val nullSafeEval = if (field.nullable) {
        s"""
         if ($row.isNullAt($ordinal)) {
           $values[$j] = null;
         } else
        """
      } else {
        ""
      }

      s"""
        final int $n = $eval.numElements();
        final Object[] $values = new Object[$n];
        for (int $j = 0; $j < $n; $j++) {
          if ($eval.isNullAt($j)) {
            $values[$j] = null;
          } else {
            final InternalRow $row = $eval.getStruct($j, $numFields);
            $nullSafeEval {
              $values[$j] = ${CodeGenerator.getValue(row, field.dataType, ordinal.toString)};
            }
          }
        }
        ${ev.value} = new $arrayClass($values);
      """
    })
  }

  override protected def withNewChildInternal(newChild: Expression): GetArrayStructFields =
    copy(child = newChild)
}

/**
 * Returns the field at `ordinal` in the Array `child`.
 *
 * We need to do type checking here as `ordinal` expression maybe unresolved.
 */
case class GetArrayItem(
    child: Expression,
    ordinal: Expression,
    failOnError: Boolean = SQLConf.get.ansiEnabled) extends BinaryExpression
  with GetArrayItemUtil
  with ExpectsInputTypes
  with ExtractValue
  with SupportQueryContext {

  override def inputTypes: Seq[AbstractDataType] = Seq(ArrayType, IntegralType)

  override def toString: String = s"$child[$ordinal]"
  override def sql: String = s"${child.sql}[${ordinal.sql}]"

  override def left: Expression = child
  override def right: Expression = ordinal
  override def nullable: Boolean =
    computeNullabilityFromArray(left, right, failOnError, nullability)
  override def dataType: DataType = child.dataType.asInstanceOf[ArrayType].elementType

  private def nullability(elements: Seq[Expression], ordinal: Int): Boolean = {
    if (ordinal >= 0 && ordinal < elements.length) {
      elements(ordinal).nullable
    } else {
      !failOnError
    }
  }

  protected override def nullSafeEval(value: Any, ordinal: Any): Any = {
    val baseValue = value.asInstanceOf[ArrayData]
    val index = ordinal.asInstanceOf[Number].intValue()
    if (index >= baseValue.numElements() || index < 0) {
      if (failOnError) {
        throw QueryExecutionErrors.invalidArrayIndexError(
          index, baseValue.numElements(), getContextOrNull())
      } else {
        null
      }
    } else if (baseValue.isNullAt(index)) {
      null
    } else {
      baseValue.get(index, dataType)
    }
  }

  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
    nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
      val index = ctx.freshName("index")
      val childArrayElementNullable = child.dataType.asInstanceOf[ArrayType].containsNull
      val nullCheck = if (childArrayElementNullable) {
        s"""else if ($eval1.isNullAt($index)) {
               ${ev.isNull} = true;
            }
         """
      } else {
        ""
      }

      val indexOutOfBoundBranch = if (failOnError) {
        val errorContext = getContextOrNullCode(ctx)
        // scalastyle:off line.size.limit
        s"throw QueryExecutionErrors.invalidArrayIndexError($index, $eval1.numElements(), $errorContext);"
        // scalastyle:on line.size.limit
      } else {
        s"${ev.isNull} = true;"
      }

      s"""
        final int $index = (int) $eval2;
        if ($index >= $eval1.numElements() || $index < 0) {
          $indexOutOfBoundBranch
        } $nullCheck else {
          ${ev.value} = ${CodeGenerator.getValue(eval1, dataType, index)};
        }
      """
    })
  }

  override protected def withNewChildrenInternal(
      newLeft: Expression, newRight: Expression): GetArrayItem =
    copy(child = newLeft, ordinal = newRight)

  override def initQueryContext(): Option[QueryContext] = if (failOnError) {
    Some(origin.context)
  } else {
    None
  }
}

/**
 * Common trait for [[GetArrayItem]] and [[ElementAt]].
 */
trait GetArrayItemUtil {

  /** `Null` is returned for invalid ordinals. */
  protected def computeNullabilityFromArray(
      child: Expression,
      ordinal: Expression,
      failOnError: Boolean,
      nullability: (Seq[Expression], Int) => Boolean): Boolean = {
    val arrayElementNullable = child.dataType.asInstanceOf[ArrayType].containsNull
    if (ordinal.foldable && !ordinal.nullable) {
      val intOrdinal = ordinal.eval().asInstanceOf[Number].intValue()
      child match {
        case CreateArray(ar, _) =>
          nullability(ar, intOrdinal)
        case GetArrayStructFields(CreateArray(elements, _), field, _, _, _) =>
          nullability(elements, intOrdinal) || field.nullable
        case _ =>
          true
      }
    } else {
      if (failOnError) arrayElementNullable || child.nullable else true
    }
  }
}

/**
 * Common trait for [[GetMapValue]] and [[ElementAt]].
 *
 * The lookup strategy is chosen once per expression instance via [[executor]]:
 *   - [[PrebuiltHashExecutor]] when the map child is foldable, its key type supports
 *     hashing, and its size meets [[SQLConf.MAP_LOOKUP_HASH_THRESHOLD]]. The hash index
 *     is built once from the constant map and reused for every row.
 *   - [[LinearExecutor]] otherwise (non-foldable maps, unsupported key types, or maps
 *     below the threshold). Hash lookup is not used here because its per-row build cost
 *     is higher than a linear scan when the index cannot be reused across rows.
 *
 * Each strategy owns both its interpreted eval and its codegen, grouped together, and
 * [[getValueEval]] / [[doGetValueGenCode]] are thin delegators.
 */
trait GetMapValueUtil extends BinaryExpression with ImplicitCastInputTypes {

  @transient private lazy val hashLookupThreshold =
    SQLConf.get.getConf(SQLConf.MAP_LOOKUP_HASH_THRESHOLD)

  @transient private lazy val executor: MapLookupExecutor = chooseExecutor()

  /**
   * True iff this expression was resolved onto the hash lookup path at construction time.
   * Visible for testing so suites can assert strategy choice without matching on
   * path-dependent inner types.
   */
  private[expressions] def usesFoldableHashLookup: Boolean =
    executor.isInstanceOf[PrebuiltHashExecutor]

  /**
   * Chooses an executor once per expression instance. The hash path is taken only for foldable
   * map inputs whose key type passes [[TypeUtils.typeWithProperEquals]] -- the same predicate
   * SPARK-55959's interpreted path already used for its hash fallback, so the interpreted and
   * codegen paths stay in sync. (SPARK-55959 also carried a codegen-only `supportsHashLookup`
   * whitelist which was a subset of this predicate; this refactor unifies both paths on one
   * predicate.)
   *
   * Concretely: non-binary-equality [[StringType]] and [[BinaryType]] keys fall through to
   * linear scan. [[VariantType]] keys are structurally unreachable here (rejected earlier by
   * [[TypeUtils.checkForMapKeyType]]). [[GeographyType]] / [[GeometryType]] are `AtomicType`s
   * with identity-inherited equals; if either becomes a valid map key in the future, proper
   * hash/equals must be defined for them -- a pre-existing constraint that already applies
   * to the interpretive hash path.
   */
  private def chooseExecutor(): MapLookupExecutor = left.dataType match {
    case MapType(keyType, _, _)
        if left.foldable && TypeUtils.typeWithProperEquals(keyType) =>
      left.eval(null) match {
        case map: MapData if map.numElements() >= hashLookupThreshold =>
          val index = buildHashIndex(map, keyType)
          val (buckets, hashMask) = buildHashBuckets(map, keyType)
          new PrebuiltHashExecutor(
            index, buckets, hashMask, map.keyArray(), map.valueArray())
        case _ => LinearExecutor
      }
    case _ => LinearExecutor
  }

  private def buildHashIndex(map: MapData, keyType: DataType): java.util.HashMap[Any, Int] = {
    val keys = map.keyArray()
    val len = keys.numElements()
    val hm = new java.util.HashMap[Any, Int]((len * 1.5).toInt)
    var i = 0
    while (i < len) {
      // putIfAbsent preserves first-match semantics for maps with duplicate keys (allowed at
      // the physical level by [[ArrayBasedMapData]]), matching the linear scan path.
      hm.putIfAbsent(keys.get(i, keyType), i)
      i += 1
    }
    hm
  }

  /**
   * Builds a power-of-two open-addressed bucket table mapping `hash(key) & hashMask` to the
   * key's index in `map.keyArray()`. Used by the codegen hash path so primitive-keyed lookups
   * avoid autoboxing. The hash function must match [[GetMapValueUtil.genHash]] for the same
   * `keyType`.
   */
  private def buildHashBuckets(map: MapData, keyType: DataType): (Array[Int], Int) = {
    val keys = map.keyArray()
    val len = keys.numElements()
    // Power-of-two capacity, load factor < 0.5, min 4. Clamped to 2^30 so (cap - 1) fits in int.
    val target = math.min(math.max(len.toLong * 2L - 1L, 1L), (1L << 30) - 1L).toInt
    val cap = math.max(java.lang.Integer.highestOneBit(target) << 1, 4)
    val buckets = new Array[Int](cap)
    java.util.Arrays.fill(buckets, -1)
    val mask = cap - 1
    var i = 0
    while (i < len) {
      var h = hashKeyOnDriver(keys.get(i, keyType), keyType) & mask
      // Open addressing with linear probing; duplicates take the next free slot so that the
      // lookup (which stops at the first match) returns the first-inserted index -- matches
      // [[buildHashIndex]] / [[ArrayBasedMapData]] first-wins semantics.
      while (buckets(h) != -1) h = (h + 1) & mask
      buckets(h) = i
      i += 1
    }
    (buckets, mask)
  }

  /** Scala-side hash for a Spark value; mirrors [[GetMapValueUtil.genHash]] per keyType. */
  private def hashKeyOnDriver(v: Any, keyType: DataType): Int = keyType match {
    case BooleanType => if (v.asInstanceOf[Boolean]) 1 else 0
    case ByteType => v.asInstanceOf[Byte].toInt
    case ShortType => v.asInstanceOf[Short].toInt
    case IntegerType | DateType | _: YearMonthIntervalType => v.asInstanceOf[Int]
    case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType | _: TimeType =>
      val l = v.asInstanceOf[Long]
      (l ^ (l >>> 32)).toInt
    case FloatType => java.lang.Float.floatToIntBits(v.asInstanceOf[Float])
    case DoubleType =>
      val l = java.lang.Double.doubleToLongBits(v.asInstanceOf[Double])
      (l ^ (l >>> 32)).toInt
    case _ => v.hashCode()
  }

  /** Java-side hash expression over the primitive/object `v`. Mirrors [[hashKeyOnDriver]]. */
  private def genHash(v: String, keyType: DataType): String = keyType match {
    case BooleanType => s"($v ? 1 : 0)"
    case ByteType | ShortType | IntegerType | DateType | _: YearMonthIntervalType => s"$v"
    case LongType | TimestampType | TimestampNTZType | _: DayTimeIntervalType | _: TimeType =>
      s"(int)($v ^ ($v >>> 32))"
    case FloatType => s"Float.floatToIntBits($v)"
    case DoubleType =>
      s"(int)(Double.doubleToLongBits($v) ^ (Double.doubleToLongBits($v) >>> 32))"
    case _ => s"$v.hashCode()"
  }

  def getValueEval(
      value: Any,
      ordinal: Any,
      keyType: DataType,
      ordering: Ordering[Any]): Any =
    executor.eval(value.asInstanceOf[MapData], ordinal, keyType, ordering)

  def doGetValueGenCode(
      ctx: CodegenContext,
      ev: ExprCode,
      mapType: MapType): ExprCode =
    executor.genCode(ctx, ev, mapType)

  /**
   * Execution strategy for a map key lookup. Inner to the trait so each concrete strategy
   * has closure access to `dataType` and `nullSafeCodeGen` on the enclosing expression.
   *
   * The `eval` signature is shaped by [[LinearExecutor]]'s needs; hash-based strategies may
   * ignore `keyType` / `ordering` (their lookup uses pre-built state that does not depend
   * on the runtime `map`, `keyType`, or `ordering`).
   */
  protected sealed trait MapLookupExecutor extends Serializable {
    def eval(map: MapData, ordinal: Any, keyType: DataType, ordering: Ordering[Any]): Any
    def genCode(ctx: CodegenContext, ev: ExprCode, mapType: MapType): ExprCode
  }

  /** Linear scan over the map's key array. Stateless; no pre-computation. */
  protected object LinearExecutor extends MapLookupExecutor {
    override def eval(
        map: MapData,
        ordinal: Any,
        keyType: DataType,
        ordering: Ordering[Any]): Any = {
      val length = map.numElements()
      val keys = map.keyArray()
      val values = map.valueArray()

      var i = 0
      var found = false
      while (i < length && !found) {
        if (ordering.equiv(keys.get(i, keyType), ordinal)) {
          found = true
        } else {
          i += 1
        }
      }

      if (!found || values.isNullAt(i)) null else values.get(i, dataType)
    }

    override def genCode(
        ctx: CodegenContext,
        ev: ExprCode,
        mapType: MapType): ExprCode = {
      val index = ctx.freshName("index")
      val length = ctx.freshName("length")
      val keys = ctx.freshName("keys")
      val values = ctx.freshName("values")
      val keyType = mapType.keyType

      val keyJavaType = CodeGenerator.javaType(keyType)
      val loopKey = ctx.freshName("loopKey")
      val i = ctx.freshName("i")

      val nullValueCheck = if (mapType.valueContainsNull) {
        s"""
           |else if ($values.isNullAt($index)) {
           |  ${ev.isNull} = true;
           |}
         """.stripMargin
      } else {
        ""
      }

      nullSafeCodeGen(ctx, ev, (eval1, eval2) => {
        s"""
           |final int $length = $eval1.numElements();
           |final ArrayData $keys = $eval1.keyArray();
           |final ArrayData $values = $eval1.valueArray();
           |int $index = -1;
           |
           |for (int $i = 0; $i < $length; $i++) {
           |  $keyJavaType $loopKey = ${CodeGenerator.getValue(keys, keyType, i)};
           |  if (${ctx.genEqual(keyType, loopKey, eval2)}) {
           |    $index = $i;
           |    break;
           |  }
           |}
           |
           |if ($index < 0) {
           |  ${ev.isNull} = true;
           |} $nullValueCheck else {
           |  ${ev.value} = ${CodeGenerator.getValue(values, dataType, index)};
           |}
         """.stripMargin
      })
    }
  }

  /**
   * Hash-based lookup for foldable (constant) maps. All state is evaluated on the driver at
   * construction time and reused for every row:
   *   - `index`: generic `HashMap[Any, Int]` used by the interpreted path; avoids duplicating
   *     the hash/equal logic on the Scala side for non-primitive keys.
   *   - `buckets` / `hashMask`: power-of-two open-addressed `int[]` table used by codegen so
   *     primitive-keyed lookups can use inline hash / primitive `==` compare without
   *     autoboxing. Populated by [[GetMapValueUtil.buildHashBuckets]] using a Scala hash that
   *     mirrors [[GetMapValueUtil.genHash]].
   *   - `keys` / `values`: the constant map's key and value arrays, embedded in the generated
   *     class via [[CodegenContext.addReferenceObj]].
   */
  protected final class PrebuiltHashExecutor(
      index: java.util.HashMap[Any, Int],
      buckets: Array[Int],
      hashMask: Int,
      keys: ArrayData,
      values: ArrayData) extends MapLookupExecutor {

    // The runtime `map`, `keyType`, and `ordering` are ignored: the map is known to be the
    // same constant on every row (the hash path is only chosen when `left.foldable`), and the
    // pre-built `index` and `values` captured at construction time are what we consult.
    override def eval(
        map: MapData,
        ordinal: Any,
        keyType: DataType,
        ordering: Ordering[Any]): Any = {
      val idx = index.getOrDefault(ordinal, -1)
      if (idx == -1 || values.isNullAt(idx)) null else values.get(idx, dataType)
    }

    override def genCode(
        ctx: CodegenContext,
        ev: ExprCode,
        mapType: MapType): ExprCode = {
      val keyType = mapType.keyType
      val bucketsRef = ctx.addReferenceObj("mapLookupBuckets", buckets, "int[]")
      val keysRef = ctx.addReferenceObj("mapLookupKeys", keys, "ArrayData")
      val valuesRef = ctx.addReferenceObj("mapLookupValues", values, "ArrayData")
      val keyJavaType = CodeGenerator.javaType(keyType)
      val h = ctx.freshName("h")
      val idx = ctx.freshName("idx")
      val candidate = ctx.freshName("candidate")
      val resultIdx = ctx.freshName("resultIdx")

      nullSafeCodeGen(ctx, ev, (_, eval2) => {
        val assignValue = if (mapType.valueContainsNull) {
          s"""
             |if ($valuesRef.isNullAt($resultIdx)) {
             |  ${ev.isNull} = true;
             |} else {
             |  ${ev.value} = ${CodeGenerator.getValue(valuesRef, dataType, resultIdx)};
             |}
           """.stripMargin
        } else {
          s"${ev.value} = ${CodeGenerator.getValue(valuesRef, dataType, resultIdx)};"
        }
        // Inline open-addressing probe. Hash must match what `buildHashBuckets` used on the
        // driver; `hashMask` is a power-of-two minus one, embedded as a literal.
        s"""
           |int $h = (${genHash(eval2, keyType)}) & $hashMask;
           |int $resultIdx = -1;
           |while ($bucketsRef[$h] != -1) {
           |  int $idx = $bucketsRef[$h];
           |  $keyJavaType $candidate = ${CodeGenerator.getValue(keysRef, keyType, idx)};
           |  if (${ctx.genEqual(keyType, candidate, eval2)}) {
           |    $resultIdx = $idx;
           |    break;
           |  }
           |  $h = ($h + 1) & $hashMask;
           |}
           |if ($resultIdx < 0) {
           |  ${ev.isNull} = true;
           |} else {
           |  $assignValue
           |}
         """.stripMargin
      })
    }
  }
}

/**
 * Returns the value of key `key` in Map `child`.
 *
 * We need to do type checking here as `key` expression maybe unresolved.
 */
case class GetMapValue(child: Expression, key: Expression)
  extends GetMapValueUtil with ExtractValue {

  @transient private lazy val ordering: Ordering[Any] =
    TypeUtils.getInterpretedOrdering(keyType)

  private[catalyst] def keyType = child.dataType.asInstanceOf[MapType].keyType

  override def checkInputDataTypes(): TypeCheckResult = child.dataType match {
    case _: MapType =>
      super.checkInputDataTypes() match {
        case f if f.isFailure => f
        case TypeCheckResult.TypeCheckSuccess =>
          TypeUtils.checkForOrderingExpr(keyType, prettyName)
      }
    // This should never happen, unless we hit a bug.
    case other => TypeCheckResult.TypeCheckFailure(
      "GetMapValue.child must be map type, but got " + other)
  }

  override def inputTypes: Seq[AbstractDataType] = Seq(MapType, keyType)

  override def toString: String = s"$child[$key]"
  override def sql: String = s"${child.sql}[${key.sql}]"

  override def left: Expression = child
  override def right: Expression = key

  /**
   * `Null` is returned for invalid ordinals.
   */
  override def nullable: Boolean = true
  override def dataType: DataType = child.dataType.asInstanceOf[MapType].valueType

  override def nullSafeEval(value: Any, ordinal: Any): Any = {
    getValueEval(value, ordinal, keyType, ordering)
  }

  override def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = {
    doGetValueGenCode(ctx, ev, child.dataType.asInstanceOf[MapType])
  }

  override protected def withNewChildrenInternal(
      newLeft: Expression, newRight: Expression): GetMapValue =
    copy(child = newLeft, key = newRight)
}
