<commit_before>/*
 * 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

import scala.collection.mutable.ArrayBuffer

import org.apache.spark.{SparkConf, SparkContext, SparkEnv, TaskContext}
import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
import org.apache.spark.internal.config
import org.apache.spark.memory.MemoryTestingUtils
import org.apache.spark.sql.catalyst.expressions.UnsafeRow
import org.apache.spark.util.collection.unsafe.sort.UnsafeExternalSorter

/**
 * Benchmark ExternalAppendOnlyUnsafeRowArray.
 * To run this benchmark:
 * {{{
 *   1. without sbt:
 *      bin/spark-submit --class <this class> --jars <spark core test jar> <spark sql test jar>
 *   2. build/sbt build/sbt ";project sql;set javaOptions
 *        in Test += \"-Dspark.memory.debugFill=false\";test:runMain <this class>"
 *   3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt ";project sql;set javaOptions
 *        in Test += \"-Dspark.memory.debugFill=false\";test:runMain <this class>"
 *      Results will be written to
 *      "benchmarks/ExternalAppendOnlyUnsafeRowArrayBenchmark-results.txt".
 * }}}
 */
object ExternalAppendOnlyUnsafeRowArrayBenchmark extends BenchmarkBase {

  private val conf = new SparkConf(false)
    // Make the Java serializer write a reset instruction (TC_RESET) after each object to test
    // for a bug we had with bytes written past the last object in a batch (SPARK-2792)
    .set(config.SERIALIZER_OBJECT_STREAM_RESET, 1)
    .set(config.SERIALIZER, "org.apache.spark.serializer.JavaSerializer")

  private def withFakeTaskContext(f: => Unit): Unit = {
    val sc = new SparkContext("local", "test", conf)
    val taskContext = MemoryTestingUtils.fakeTaskContext(SparkEnv.get)
    TaskContext.setTaskContext(taskContext)
    f
    sc.stop()
  }

  private def testRows(numRows: Int): Seq[UnsafeRow] = {
    val random = new java.util.Random()
    (1 to numRows).map(_ => {
      val row = new UnsafeRow(1)
      row.pointTo(new Array[Byte](64), 16)
      row.setLong(0, random.nextLong())
      row
    })
  }

  def testAgainstRawArrayBuffer(numSpillThreshold: Int, numRows: Int, iterations: Int): Unit = {
    val rows = testRows(numRows)

    val benchmark = new Benchmark(s"Array with $numRows rows", iterations * numRows,
      output = output)
<commit_msg>[SPARK-34832][SQL][TEST] Set EXECUTOR_ALLOW_SPARK_CONTEXT to true to ensure ExternalAppendOnlyUnsafeRowArrayBenchmark run successfully<commit_after>/*
 * 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

import scala.collection.mutable.ArrayBuffer

import org.apache.spark.{SparkConf, SparkContext, SparkEnv, TaskContext}
import org.apache.spark.benchmark.{Benchmark, BenchmarkBase}
import org.apache.spark.internal.config
import org.apache.spark.memory.MemoryTestingUtils
import org.apache.spark.sql.catalyst.expressions.UnsafeRow
import org.apache.spark.util.collection.unsafe.sort.UnsafeExternalSorter

/**
 * Benchmark ExternalAppendOnlyUnsafeRowArray.
 * To run this benchmark:
 * {{{
 *   1. without sbt:
 *      bin/spark-submit --class <this class> --jars <spark core test jar> <spark sql test jar>
 *   2. build/sbt build/sbt ";project sql;set javaOptions
 *        in Test += \"-Dspark.memory.debugFill=false\";test:runMain <this class>"
 *   3. generate result: SPARK_GENERATE_BENCHMARK_FILES=1 build/sbt ";project sql;set javaOptions
 *        in Test += \"-Dspark.memory.debugFill=false\";test:runMain <this class>"
 *      Results will be written to
 *      "benchmarks/ExternalAppendOnlyUnsafeRowArrayBenchmark-results.txt".
 * }}}
 */
object ExternalAppendOnlyUnsafeRowArrayBenchmark extends BenchmarkBase {

  private val conf = new SparkConf(false)
    // Make the Java serializer write a reset instruction (TC_RESET) after each object to test
    // for a bug we had with bytes written past the last object in a batch (SPARK-2792)
    .set(config.SERIALIZER_OBJECT_STREAM_RESET, 1)
    .set(config.SERIALIZER, "org.apache.spark.serializer.JavaSerializer")
    // SPARK-34832: Add this configuration to allow `withFakeTaskContext` method
    // to create `SparkContext` on the executor side.
    .set(config.EXECUTOR_ALLOW_SPARK_CONTEXT, true)

  private def withFakeTaskContext(f: => Unit): Unit = {
    val sc = new SparkContext("local", "test", conf)
    val taskContext = MemoryTestingUtils.fakeTaskContext(SparkEnv.get)
    TaskContext.setTaskContext(taskContext)
    f
    sc.stop()
  }

  private def testRows(numRows: Int): Seq[UnsafeRow] = {
    val random = new java.util.Random()
    (1 to numRows).map(_ => {
      val row = new UnsafeRow(1)
      row.pointTo(new Array[Byte](64), 16)
      row.setLong(0, random.nextLong())
      row
    })
  }

  def testAgainstRawArrayBuffer(numSpillThreshold: Int, numRows: Int, iterations: Int): Unit = {
    val rows = testRows(numRows)

    val benchmark = new Benchmark(s"Array with $numRows rows", iterations * numRows,
      output = output)

{'fraction_non_alphanumeric': 0.09652093596059114, 'fraction_numerical': 0.00769704433497537, 'mean_word_length': 4.0053929121725735, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '6277976', 'n_tokens_mistral': 1976, 'n_tokens_neox': 1880, 'n_words': 604}